{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1b79ff35-50a3-40cd-86d9-703f1f8cd2c5",
   "metadata": {},
   "source": [
    "# Citations\n",
    "\n",
    "How can we get a model to cite which parts of the source documents it referenced in its response?\n",
    "\n",
    "To explore some techniques for extracting citations, let's first create a simple RAG chain. To start we'll just retrieve from Wikipedia using the [WikipediaRetriever](https://api.python.langchain.com/en/latest/retrievers/langchain_community.retrievers.wikipedia.WikipediaRetriever.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a70c423-f61f-4230-b70a-d3605b31afab",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "First we'll need to install some dependencies and set environment vars for the models we'll be using."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f1d26ded-e8d5-4f80-86b9-26d464869175",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -qU langchain langchain-openai langchain-anthropic langchain-community wikipedia"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8732a85a-dd1a-483c-8da7-a81251276aa1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = getpass.getpass()\n",
    "os.environ[\"ANTHROPIC_API_KEY\"] = getpass.getpass()\n",
    "\n",
    "# Uncomment if you want to log to LangSmith\n",
    "# os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\n",
    "# os.environ[\"LANGCHAIN_API_KEY\"] = getpass.getpass()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4e17c3f6-8ce6-4767-b615-50a57c84c7b0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m System Message \u001b[0m================================\n",
      "\n",
      "You're a helpful AI assistant. Given a user question and some Wikipedia article snippets, answer the user question. If none of the articles answer the question, just say you don't know.\n",
      "\n",
      "Here are the Wikipedia articles:\u001b[33;1m\u001b[1;3m{context}\u001b[0m\n",
      "\n",
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "\u001b[33;1m\u001b[1;3m{question}\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "from langchain_community.retrievers import WikipediaRetriever\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "llm = ChatOpenAI(model=\"gpt-3.5-turbo\", temperature=0)\n",
    "wiki = WikipediaRetriever(top_k_results=6, doc_content_chars_max=2000)\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            \"You're a helpful AI assistant. Given a user question and some Wikipedia article snippets, answer the user question. If none of the articles answer the question, just say you don't know.\\n\\nHere are the Wikipedia articles:{context}\",\n",
    "        ),\n",
    "        (\"human\", \"{question}\"),\n",
    "    ]\n",
    ")\n",
    "prompt.pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c89e2045-9244-43e6-bf3f-59af22658529",
   "metadata": {},
   "source": [
    "Now that we've got a model, retriver and prompt, let's chain them all together. We'll need to add some logic for formatting our retrieved Documents to a string that can be passed to our prompt. We'll make it so our chain returns both the answer and the retrieved Documents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e6b55c52-9062-4061-9017-d8a1bce72078",
   "metadata": {},
   "outputs": [],
   "source": [
    "from operator import itemgetter\n",
    "from typing import List\n",
    "\n",
    "from langchain_core.documents import Document\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.runnables import (\n",
    "    RunnableLambda,\n",
    "    RunnableParallel,\n",
    "    RunnablePassthrough,\n",
    ")\n",
    "\n",
    "\n",
    "def format_docs(docs: List[Document]) -> str:\n",
    "    \"\"\"Convert Documents to a single string.:\"\"\"\n",
    "    formatted = [\n",
    "        f\"Article Title: {doc.metadata['title']}\\nArticle Snippet: {doc.page_content}\"\n",
    "        for doc in docs\n",
    "    ]\n",
    "    return \"\\n\\n\" + \"\\n\\n\".join(formatted)\n",
    "\n",
    "\n",
    "format = itemgetter(\"docs\") | RunnableLambda(format_docs)\n",
    "# subchain for generating an answer once we've done retrieval\n",
    "answer = prompt | llm | StrOutputParser()\n",
    "# complete chain that calls wiki -> formats docs to string -> runs answer subchain -> returns just the answer and retrieved docs.\n",
    "chain = (\n",
    "    RunnableParallel(question=RunnablePassthrough(), docs=wiki)\n",
    "    .assign(context=format)\n",
    "    .assign(answer=answer)\n",
    "    .pick([\"answer\", \"docs\"])\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "29568bb9-a9b0-4b63-9884-79fd88279ac6",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'answer': 'Cheetahs are capable of running at speeds between 93 to 104 km/h (58 to 65 mph). They have evolved specialized adaptations for speed, including a light build, long thin legs, and a long tail.',\n",
       " 'docs': [Document(page_content='The cheetah (Acinonyx jubatus) is a large cat and the fastest land animal. It has a tawny to creamy white or pale buff fur that is marked with evenly spaced, solid black spots. The head is small and rounded, with a short snout and black tear-like facial streaks. It reaches 67–94 cm (26–37 in) at the shoulder, and the head-and-body length is between 1.1 and 1.5 m (3 ft 7 in and 4 ft 11 in). Adults weigh between 21 and 72 kg (46 and 159 lb). The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail.\\nThe cheetah was first described in the late 18th century. Four subspecies are recognised today that are native to Africa and central Iran. An African subspecies was introduced to India in 2022. It is now distributed mainly in small, fragmented populations in northwestern, eastern and southern Africa and central Iran. It lives in a variety of habitats such as savannahs in the Serengeti, arid mountain ranges in the Sahara, and hilly desert terrain.\\nThe cheetah lives in three main social groups: females and their cubs, male \"coalitions\", and solitary males. While females lead a nomadic life searching for prey in large home ranges, males are more sedentary and instead establish much smaller territories in areas with plentiful prey and access to females. The cheetah is active during the day, with peaks during dawn and dusk. It feeds on small- to medium-sized prey, mostly weighing under 40 kg (88 lb), and prefers medium-sized ungulates such as impala, springbok and Thomson\\'s gazelles. The cheetah typically stalks its prey within 60–100 m (200–330 ft) before charging towards it, trips it during the chase and bites its throat to suffocate it to death. It breeds throughout the year. After a gestation of nearly three months, females give birth to a litter of three or four cubs. Cheetah cubs are highly vulnerable to predation by other large carnivores. They are weaned a', metadata={'title': 'Cheetah', 'summary': 'The cheetah (Acinonyx jubatus) is a large cat and the fastest land animal. It has a tawny to creamy white or pale buff fur that is marked with evenly spaced, solid black spots. The head is small and rounded, with a short snout and black tear-like facial streaks. It reaches 67–94 cm (26–37 in) at the shoulder, and the head-and-body length is between 1.1 and 1.5 m (3 ft 7 in and 4 ft 11 in). Adults weigh between 21 and 72 kg (46 and 159 lb). The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail.\\nThe cheetah was first described in the late 18th century. Four subspecies are recognised today that are native to Africa and central Iran. An African subspecies was introduced to India in 2022. It is now distributed mainly in small, fragmented populations in northwestern, eastern and southern Africa and central Iran. It lives in a variety of habitats such as savannahs in the Serengeti, arid mountain ranges in the Sahara, and hilly desert terrain.\\nThe cheetah lives in three main social groups: females and their cubs, male \"coalitions\", and solitary males. While females lead a nomadic life searching for prey in large home ranges, males are more sedentary and instead establish much smaller territories in areas with plentiful prey and access to females. The cheetah is active during the day, with peaks during dawn and dusk. It feeds on small- to medium-sized prey, mostly weighing under 40 kg (88 lb), and prefers medium-sized ungulates such as impala, springbok and Thomson\\'s gazelles. The cheetah typically stalks its prey within 60–100 m (200–330 ft) before charging towards it, trips it during the chase and bites its throat to suffocate it to death. It breeds throughout the year. After a gestation of nearly three months, females give birth to a litter of three or four cubs. Cheetah cubs are highly vulnerable to predation by other large carnivores. They are weaned at around four months and are independent by around 20 months of age.\\nThe cheetah is threatened by habitat loss, conflict with humans, poaching and high susceptibility to diseases. In 2016, the global cheetah population was estimated at 7,100 individuals in the wild; it is listed as Vulnerable on the IUCN Red List. It has been widely depicted in art, literature, advertising, and animation. It was tamed in ancient Egypt and trained for hunting ungulates in the Arabian Peninsula and India. It has been kept in zoos since the early 19th century.\\n\\n', 'source': 'https://en.wikipedia.org/wiki/Cheetah'}),\n",
       "  Document(page_content='More than 70 years after India\\'s native subspecies of the cheetah—the Asiatic cheetah (Acinonyx jubatus venaticus)—became extinct there, small numbers of Southeast African cheetah (Acinonyx jubatus jubatus) have been flown in from Namibia and South Africa to a national park in India.  The experiment has been permitted by India\\'s supreme court on a short-term basis to test long-term adaptation. The Asiatic subspecies is now found only in Iran in critically endangered numbers.The Asiatic cheetah whose long history on the Indian subcontinent gave the Sanskrit-derived vernacular name \"cheetah\", or \"spotted\", to the entire species,  Acinonyx jubatus, also had a gradual history of habitat loss there. In Punjab, before the thorn forests were cleared for agriculture and human settlement,  they were intermixed with open grasslands grazed by large herds of blackbuck; these co-existed with their main natural predator, the Asiatic cheetah.  The blackbuck is no longer extant in Punjab. Later, more habitat loss, prey depletion, and trophy hunting were to lead to the extinction of the Asiatic cheetah in other regions of India.\\nDiscussions on cheetah reintroduction in India began soon after extinction was confirmed, in the mid-1950s. Proposals were made to the governments of Iran from the 1970s, but fell through chiefly for reasons of political instability there. Offers from Kenya for introducing African cheetahs were made as early as the 1980s. Proposals for the introduction of African cheetahs were made by the Indian government in 2009, but disallowed by India\\'s supreme court. The court reversed its decision in early 2020, allowing the import of a small number, on an experimental basis for testing long-term adaptation. On 17 September 2022, five female and three male southeast African cheetahs, between the ages of four and six (a gift from the government of Namibia), were released in a small quarantined enclosure within the Kuno National Park in the state of Madhya Pradesh. The c', metadata={'title': 'Cheetah reintroduction in India', 'summary': 'More than 70 years after India\\'s native subspecies of the cheetah—the Asiatic cheetah (Acinonyx jubatus venaticus)—became extinct there, small numbers of Southeast African cheetah (Acinonyx jubatus jubatus) have been flown in from Namibia and South Africa to a national park in India.  The experiment has been permitted by India\\'s supreme court on a short-term basis to test long-term adaptation. The Asiatic subspecies is now found only in Iran in critically endangered numbers.The Asiatic cheetah whose long history on the Indian subcontinent gave the Sanskrit-derived vernacular name \"cheetah\", or \"spotted\", to the entire species,  Acinonyx jubatus, also had a gradual history of habitat loss there. In Punjab, before the thorn forests were cleared for agriculture and human settlement,  they were intermixed with open grasslands grazed by large herds of blackbuck; these co-existed with their main natural predator, the Asiatic cheetah.  The blackbuck is no longer extant in Punjab. Later, more habitat loss, prey depletion, and trophy hunting were to lead to the extinction of the Asiatic cheetah in other regions of India.\\nDiscussions on cheetah reintroduction in India began soon after extinction was confirmed, in the mid-1950s. Proposals were made to the governments of Iran from the 1970s, but fell through chiefly for reasons of political instability there. Offers from Kenya for introducing African cheetahs were made as early as the 1980s. Proposals for the introduction of African cheetahs were made by the Indian government in 2009, but disallowed by India\\'s supreme court. The court reversed its decision in early 2020, allowing the import of a small number, on an experimental basis for testing long-term adaptation. On 17 September 2022, five female and three male southeast African cheetahs, between the ages of four and six (a gift from the government of Namibia), were released in a small quarantined enclosure within the Kuno National Park in the state of Madhya Pradesh. The cheetahs, all fitted with radio collars, will remain in the quarantined enclosure for a month; initially, the males (and later the females) will be released into the 748.76 km2 (289.10 sq mi) park. The relocation has been supervised by Yadvendradev V. Jhala of the Wildlife Institute of India and zoologist Laurie Marker, of the Namibia-based Cheetah Conservation Fund. Subsequently, 12 cheetahs from South Africa will be released in Kuno;   eventually, the total number of African cheetahs in Kuno will be brought up to 40 individuals.  As of Jan 16, 2024, seven adult cheetahs from Africa and three cubs (of four born in Kuno two months earlier) had died in Kuno National Park.\\nThe scientific reaction to the translocation has been mixed. Adrian Tordiffe (a wildlife veterinary pharmacologist at the University of Pretoria who will be supervising the release of the cheetahs) is an enthusiast,  who views India as providing \"protected space\" for the fragmented and threatened population of the world\\'s cheetahs.  K. Ullas Karanth, one of India\\'s tiger experts, has been critical of the effort, considering it to be a \"PR exercise.\" India\\'s \"realities\", he says, such as human overpopulation, and the presence of larger feline predators and packs of feral dogs, could all cause potentially \"high mortalities,\" and require a continual import of African cheetahs. Kuno National Park is a relatively new national park, having received that status in 2018. It had been founded previously as a wildlife sanctuary to implement the Asiatic Lion Reintroduction Project, which aimed to establish a second Asiatic lion population in India. The goal was to protect the isolated lions of the Gir National Park (in Gujarat) from a potential mass mortality event, set off by the outbreak of an epizootic. Although the state government of Gujarat was ordered by India\\'s Supreme Court in April 2013 to transfer a small population of lions from Gujarat to Kuno, and was given six months to complete the transfer, they ultimately  resisted implementing the order.', 'source': 'https://en.wikipedia.org/wiki/Cheetah_reintroduction_in_India'}),\n",
       "  Document(page_content='The Southeast African cheetah (Acinonyx jubatus jubatus) is the nominate cheetah subspecies native to East and Southern Africa. The Southern African cheetah lives mainly in the lowland areas and deserts of the Kalahari, the savannahs of Okavango Delta, and the grasslands of the Transvaal region in South Africa. In Namibia, cheetahs are mostly found in farmlands. In India, four cheetahs of the subspecies are living in Kuno National Park in Madhya Pradesh after having been introduced there.\\n\\n\\n== Taxonomy ==\\nThe Southern African cheetah was first described by German naturalist Johann Christian Daniel von Schreber in his book Die Säugethiere in Abbildungen nach der Natur mit Beschreibungen (The Mammals illustrated as in Nature with Descriptions), published in 1775. Schreber described the species on basis of a specimen from the Cape of Good Hope. It is therefore the nominate subspecies. Subpopulations have been called \"South African cheetah\" and \"Namibian cheetah.\"Following Schreber\\'s description, other naturalists and zoologists also described cheetah specimens from many parts of Southern and East Africa that today are all considered synonyms of A. j. jubatus:\\nFelis guttata proposed in 1804 by Johann Hermann;\\nFelis fearonii proposed in 1834 by Andrew Smith;\\nFelis lanea proposed in 1877 by Philip Sclater;\\nAcinonyx jubatus obergi proposed in 1913 by Max Hilzheimer;\\nAcinonyx jubatus ngorongorensis proposed in 1913 by Hilzheimer on basis of a specimen from Ngorongoro, German East Africa;\\nAcinonyx jubatus velox proposed in 1913 by Edmund Heller on basis of a cheetah that was shot by Kermit Roosevelt in June 1909 in the Kenyan highlands.\\nAcinonyx rex proposed in 1927 by Reginald Innes Pocock on basis of a specimen from the Umvukwe Range in Rhodesia.In 2005, the authors of Mammal Species of the World grouped A. j. guttata, A. j. lanea, A. j. obergi, and A. j. rex under A j. jubatus, whilst recognizing A. j. raineyi and A. j. velox as valid taxa and considering P. l. ngorongore', metadata={'title': 'Southeast African cheetah', 'summary': 'The Southeast African cheetah (Acinonyx jubatus jubatus) is the nominate cheetah subspecies native to East and Southern Africa. The Southern African cheetah lives mainly in the lowland areas and deserts of the Kalahari, the savannahs of Okavango Delta, and the grasslands of the Transvaal region in South Africa. In Namibia, cheetahs are mostly found in farmlands. In India, four cheetahs of the subspecies are living in Kuno National Park in Madhya Pradesh after having been introduced there.', 'source': 'https://en.wikipedia.org/wiki/Southeast_African_cheetah'}),\n",
       "  Document(page_content='Footspeed, or sprint speed, is the maximum speed at which a human can run. It is affected by many factors, varies greatly throughout the population, and is important in athletics and many sports, such as association football, rugby football, American football,  track and field, field hockey, tennis, baseball, and basketball.\\n\\n\\n== Factors in speed ==\\nThe key determinant of footspeed in sprinting is the predominance of one distinct type of muscle fibre over another, specifically the ratio of fast-twitch muscles to slow-twitch muscles in a sprinter\\'s physical makeup.  Though fast-twitch muscles produce no more energy than slow-twitch muscles when they contract, they do so more rapidly through a process of anaerobic metabolism, though at the cost of inferior efficiency over longer periods of firing.  The average human has an almost-equal ratio of fast-twitch to slow-twitch fibers, but top sprinters may have as much as 80% fast-twitch fibers, while top long-distance runners may have only 20%. This ratio is believed to have genetic origins, though some assert that it can be adjusted by muscle training.  \"Speed camps\" and \"Speed Training Manuals\", which purport to provide fractional increases in maximum footspeed, are popular among budding professional athletes, and some sources estimate that 17–19% of speed can be trained.Though good running form is useful in increasing speed, fast and slow runners have been shown to move their legs at nearly the same rate – it is the force exerted by the leg on the ground that separates fast sprinters from slow.  Top short-distance runners exert as much as four times their body weight in pressure on the running surface. For this reason, muscle mass in the legs, relative to total body weight, is a key factor in maximizing footspeed.\\n\\n\\n== Limits of speed ==\\nThe record is 44.72 km/h (27.78 mph), measured between meter 60 and meter 80 of the 100 meters sprint at the 2009 World Championships in Athletics by Usain Bolt. (Bolt\\'s average speed o', metadata={'title': 'Footspeed', 'summary': 'Footspeed, or sprint speed, is the maximum speed at which a human can run. It is affected by many factors, varies greatly throughout the population, and is important in athletics and many sports, such as association football, rugby football, American football,  track and field, field hockey, tennis, baseball, and basketball.\\n\\n', 'source': 'https://en.wikipedia.org/wiki/Footspeed'}),\n",
       "  Document(page_content=\"This is a list of the fastest animals in the world, by types of animal.\\n\\n\\n== Fastest organism ==\\nThe peregrine falcon is the fastest bird, and the fastest member of the animal kingdom, with a diving speed of over 300 km/h (190 mph). The fastest land animal is the cheetah. Among the fastest animals in the sea is the black marlin, with uncertain and conflicting reports of recorded speeds.When drawing comparisons between different classes of animals, an alternative unit is sometimes used for organisms: body length per second. On this basis the 'fastest' organism on earth, relative to its body length, is the Southern Californian mite, Paratarsotomus macropalpis, which has a speed of 322 body lengths per second. The equivalent speed for a human, running as fast as this mite, would be 1,300 mph (2,092 km/h), or approximately Mach 1.7. The speed of the P. macropalpis is far in excess of the previous record holder, the Australian tiger beetle Cicindela eburneola, which is the fastest insect in the world relative to body size, with a recorded speed of 1.86 metres per second (6.7 km/h; 4.2 mph), or 171 body lengths per second. The cheetah, the fastest land mammal, scores at only 16 body lengths per second, while Anna's hummingbird has the highest known length-specific velocity attained by any vertebrate.\\n\\n\\n== Invertebrates ==\\n\\n\\n== Fish ==\\nDue to physical constraints, fish may be incapable of exceeding swim speeds of 36 km/h (22 mph). The larger reported figures below are therefore highly questionable:\\n\\n\\n== Amphibians ==\\n\\n\\n== Reptiles ==\\n\\n\\n== Birds ==\\n\\n\\n== Mammals ==\\n\\n\\n== See also ==\\nSpeed records\\n\\n\\n== Notes ==\\n\\n\\n== References ==\", metadata={'title': 'Fastest animals', 'summary': 'This is a list of the fastest animals in the world, by types of animal.', 'source': 'https://en.wikipedia.org/wiki/Fastest_animals'}),\n",
       "  Document(page_content=\"Pursuit predation is a form of predation in which predators actively give chase to their prey, either solitarily or as a group.  It is an alternate predation strategy to ambush predation — pursuit predators rely on superior speed, endurance and/or teamwork to seize the prey, while ambush predators use concealment, luring, exploiting of surroundings and the element of surprise to capture the prey.  While the two patterns of predation are not mutually exclusive, morphological differences in an organism's body plan can create an evolutionary bias favoring either type of predation.\\nPursuit predation is typically observed in carnivorous species within the kingdom Animalia, such as cheetahs, lions, wolves and early Homo species.  The chase can be initiated either by the predator, or by the prey if it is alerted to a predator's presence and attempt to flee before the predator gets close.  The chase ends either when the predator successfully catches up and tackles the prey, or when the predator abandons the attempt after the prey outruns it and escapes.\\nOne particular form of pursuit predation is persistence hunting, where the predator stalks the prey slowly but persistently to wear it down physically with fatigue or overheating; some animals are examples of both types of pursuit.\\n\\n\\n== Strategy ==\\nThere is still uncertainty as to whether predators behave with a general tactic or strategy while preying. However, among pursuit predators there are several common behaviors. Often, predators will scout potential prey, assessing prey quantity and density prior to engaging in a pursuit. Certain predators choose to pursue prey primarily in a group of conspecifics; these animals are known as pack hunters or group pursuers. Other species choose to hunt alone. These two behaviors are typically due to differences in hunting success, where some groups are very successful in groups and others are more successful alone. Pursuit predators may also choose to either exhaust their metabolic r\", metadata={'title': 'Pursuit predation', 'summary': \"Pursuit predation is a form of predation in which predators actively give chase to their prey, either solitarily or as a group.  It is an alternate predation strategy to ambush predation — pursuit predators rely on superior speed, endurance and/or teamwork to seize the prey, while ambush predators use concealment, luring, exploiting of surroundings and the element of surprise to capture the prey.  While the two patterns of predation are not mutually exclusive, morphological differences in an organism's body plan can create an evolutionary bias favoring either type of predation.\\nPursuit predation is typically observed in carnivorous species within the kingdom Animalia, such as cheetahs, lions, wolves and early Homo species.  The chase can be initiated either by the predator, or by the prey if it is alerted to a predator's presence and attempt to flee before the predator gets close.  The chase ends either when the predator successfully catches up and tackles the prey, or when the predator abandons the attempt after the prey outruns it and escapes.\\nOne particular form of pursuit predation is persistence hunting, where the predator stalks the prey slowly but persistently to wear it down physically with fatigue or overheating; some animals are examples of both types of pursuit.\", 'source': 'https://en.wikipedia.org/wiki/Pursuit_predation'})]}"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.invoke(\"How fast are cheetahs?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f1f9a49-8f3f-44dd-98df-0218b5fb93a6",
   "metadata": {},
   "source": [
    "LangSmith trace: https://smith.langchain.com/public/4bc9a13a-d320-46dc-a70c-7109641e7308/r"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7619ba1-33bd-48bf-8637-be409c94037f",
   "metadata": {},
   "source": [
    "## Function-calling\n",
    "\n",
    "### Cite documents\n",
    "Let's try using [OpenAI function-calling](/docs/modules/model_io/chat/function_calling) to make the model specify which of the provided documents it's actually referencing when answering. LangChain has some utils for converting Pydantic ojbects to the JSONSchema format expected by OpenAI, so we'll use that to define our functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0af2c3a1-870c-428e-95da-0c2fd04d5616",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.pydantic_v1 import BaseModel, Field\n",
    "\n",
    "\n",
    "class cited_answer(BaseModel):\n",
    "    \"\"\"Answer the user question based only on the given sources, and cite the sources used.\"\"\"\n",
    "\n",
    "    answer: str = Field(\n",
    "        ...,\n",
    "        description=\"The answer to the user question, which is based only on the given sources.\",\n",
    "    )\n",
    "    citations: List[int] = Field(\n",
    "        ...,\n",
    "        description=\"The integer IDs of the SPECIFIC sources which justify the answer.\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68b95186-faf5-46f1-8715-ebbc38207d5d",
   "metadata": {},
   "source": [
    "Let's see what the model output is like when we pass in our functions and a user input:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2770391d-c656-43cb-95f5-2098b8c6d706",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_0VO8uyUo16jzq86FQDoka2zQ', 'function': {'arguments': '{\\n  \"answer\": \"Brian\\'s height is 6\\'2\\\\\" - 3 inches\",\\n  \"citations\": [1, 3]\\n}', 'name': 'cited_answer'}, 'type': 'function'}]})"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "llm_with_tool = llm.bind_tools(\n",
    "    [cited_answer],\n",
    "    tool_choice=\"cited_answer\",\n",
    ")\n",
    "example_q = \"\"\"What Brian's height?\n",
    "\n",
    "Source: 1\n",
    "Information: Suzy is 6'2\"\n",
    "\n",
    "Source: 2\n",
    "Information: Jeremiah is blonde\n",
    "\n",
    "Source: 3\n",
    "Information: Brian is 3 inches shorted than Suzy\"\"\"\n",
    "llm_with_tool.invoke(example_q)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b847b53-987e-4d3a-9621-77e613d49cfd",
   "metadata": {},
   "source": [
    "We'll add an output parser to convert the OpenAI API response to a nice dictionary. We use the [JsonOutputKeyToolsParser](https://api.python.langchain.com/en/latest/output_parsers/langchain.output_parsers.openai_tools.JsonOutputKeyToolsParser.html#langchain.output_parsers.openai_tools.JsonOutputKeyToolsParser) for this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "bcf7830e-099f-4ea9-8636-560ebbd77692",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'answer': 'Brian\\'s height is 6\\'2\" - 3 inches', 'citations': [1, 3]}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.output_parsers.openai_tools import JsonOutputKeyToolsParser\n",
    "\n",
    "output_parser = JsonOutputKeyToolsParser(key_name=\"cited_answer\", return_single=True)\n",
    "(llm_with_tool | output_parser).invoke(example_q)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb8bbbb5-2afc-401f-a140-648c3d2c4522",
   "metadata": {},
   "source": [
    "Now we're ready to put together our chain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "3cb835f3-3cf5-4144-bf6b-24558b9faf31",
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_docs_with_id(docs: List[Document]) -> str:\n",
    "    formatted = [\n",
    "        f\"Source ID: {i}\\nArticle Title: {doc.metadata['title']}\\nArticle Snippet: {doc.page_content}\"\n",
    "        for i, doc in enumerate(docs)\n",
    "    ]\n",
    "    return \"\\n\\n\" + \"\\n\\n\".join(formatted)\n",
    "\n",
    "\n",
    "format_1 = itemgetter(\"docs\") | RunnableLambda(format_docs_with_id)\n",
    "answer_1 = prompt | llm_with_tool | output_parser\n",
    "chain_1 = (\n",
    "    RunnableParallel(question=RunnablePassthrough(), docs=wiki)\n",
    "    .assign(context=format_1)\n",
    "    .assign(cited_answer=answer_1)\n",
    "    .pick([\"cited_answer\", \"docs\"])\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "3e259b2f-5147-4c3c-9c26-b4eb8143e5f0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'cited_answer': {'answer': 'Cheetahs can run at speeds of 93 to 104 km/h (58 to 65 mph).',\n",
       "  'citations': [0]},\n",
       " 'docs': [Document(page_content='The cheetah (Acinonyx jubatus) is a large cat and the fastest land animal. It has a tawny to creamy white or pale buff fur that is marked with evenly spaced, solid black spots. The head is small and rounded, with a short snout and black tear-like facial streaks. It reaches 67–94 cm (26–37 in) at the shoulder, and the head-and-body length is between 1.1 and 1.5 m (3 ft 7 in and 4 ft 11 in). Adults weigh between 21 and 72 kg (46 and 159 lb). The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail.\\nThe cheetah was first described in the late 18th century. Four subspecies are recognised today that are native to Africa and central Iran. An African subspecies was introduced to India in 2022. It is now distributed mainly in small, fragmented populations in northwestern, eastern and southern Africa and central Iran. It lives in a variety of habitats such as savannahs in the Serengeti, arid mountain ranges in the Sahara, and hilly desert terrain.\\nThe cheetah lives in three main social groups: females and their cubs, male \"coalitions\", and solitary males. While females lead a nomadic life searching for prey in large home ranges, males are more sedentary and instead establish much smaller territories in areas with plentiful prey and access to females. The cheetah is active during the day, with peaks during dawn and dusk. It feeds on small- to medium-sized prey, mostly weighing under 40 kg (88 lb), and prefers medium-sized ungulates such as impala, springbok and Thomson\\'s gazelles. The cheetah typically stalks its prey within 60–100 m (200–330 ft) before charging towards it, trips it during the chase and bites its throat to suffocate it to death. It breeds throughout the year. After a gestation of nearly three months, females give birth to a litter of three or four cubs. Cheetah cubs are highly vulnerable to predation by other large carnivores. They are weaned a', metadata={'title': 'Cheetah', 'summary': 'The cheetah (Acinonyx jubatus) is a large cat and the fastest land animal. It has a tawny to creamy white or pale buff fur that is marked with evenly spaced, solid black spots. The head is small and rounded, with a short snout and black tear-like facial streaks. It reaches 67–94 cm (26–37 in) at the shoulder, and the head-and-body length is between 1.1 and 1.5 m (3 ft 7 in and 4 ft 11 in). Adults weigh between 21 and 72 kg (46 and 159 lb). The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail.\\nThe cheetah was first described in the late 18th century. Four subspecies are recognised today that are native to Africa and central Iran. An African subspecies was introduced to India in 2022. It is now distributed mainly in small, fragmented populations in northwestern, eastern and southern Africa and central Iran. It lives in a variety of habitats such as savannahs in the Serengeti, arid mountain ranges in the Sahara, and hilly desert terrain.\\nThe cheetah lives in three main social groups: females and their cubs, male \"coalitions\", and solitary males. While females lead a nomadic life searching for prey in large home ranges, males are more sedentary and instead establish much smaller territories in areas with plentiful prey and access to females. The cheetah is active during the day, with peaks during dawn and dusk. It feeds on small- to medium-sized prey, mostly weighing under 40 kg (88 lb), and prefers medium-sized ungulates such as impala, springbok and Thomson\\'s gazelles. The cheetah typically stalks its prey within 60–100 m (200–330 ft) before charging towards it, trips it during the chase and bites its throat to suffocate it to death. It breeds throughout the year. After a gestation of nearly three months, females give birth to a litter of three or four cubs. Cheetah cubs are highly vulnerable to predation by other large carnivores. They are weaned at around four months and are independent by around 20 months of age.\\nThe cheetah is threatened by habitat loss, conflict with humans, poaching and high susceptibility to diseases. In 2016, the global cheetah population was estimated at 7,100 individuals in the wild; it is listed as Vulnerable on the IUCN Red List. It has been widely depicted in art, literature, advertising, and animation. It was tamed in ancient Egypt and trained for hunting ungulates in the Arabian Peninsula and India. It has been kept in zoos since the early 19th century.\\n\\n', 'source': 'https://en.wikipedia.org/wiki/Cheetah'}),\n",
       "  Document(page_content='More than 70 years after India\\'s native subspecies of the cheetah—the Asiatic cheetah (Acinonyx jubatus venaticus)—became extinct there, small numbers of Southeast African cheetah (Acinonyx jubatus jubatus) have been flown in from Namibia and South Africa to a national park in India.  The experiment has been permitted by India\\'s supreme court on a short-term basis to test long-term adaptation. The Asiatic subspecies is now found only in Iran in critically endangered numbers.The Asiatic cheetah whose long history on the Indian subcontinent gave the Sanskrit-derived vernacular name \"cheetah\", or \"spotted\", to the entire species,  Acinonyx jubatus, also had a gradual history of habitat loss there. In Punjab, before the thorn forests were cleared for agriculture and human settlement,  they were intermixed with open grasslands grazed by large herds of blackbuck; these co-existed with their main natural predator, the Asiatic cheetah.  The blackbuck is no longer extant in Punjab. Later, more habitat loss, prey depletion, and trophy hunting were to lead to the extinction of the Asiatic cheetah in other regions of India.\\nDiscussions on cheetah reintroduction in India began soon after extinction was confirmed, in the mid-1950s. Proposals were made to the governments of Iran from the 1970s, but fell through chiefly for reasons of political instability there. Offers from Kenya for introducing African cheetahs were made as early as the 1980s. Proposals for the introduction of African cheetahs were made by the Indian government in 2009, but disallowed by India\\'s supreme court. The court reversed its decision in early 2020, allowing the import of a small number, on an experimental basis for testing long-term adaptation. On 17 September 2022, five female and three male southeast African cheetahs, between the ages of four and six (a gift from the government of Namibia), were released in a small quarantined enclosure within the Kuno National Park in the state of Madhya Pradesh. The c', metadata={'title': 'Cheetah reintroduction in India', 'summary': 'More than 70 years after India\\'s native subspecies of the cheetah—the Asiatic cheetah (Acinonyx jubatus venaticus)—became extinct there, small numbers of Southeast African cheetah (Acinonyx jubatus jubatus) have been flown in from Namibia and South Africa to a national park in India.  The experiment has been permitted by India\\'s supreme court on a short-term basis to test long-term adaptation. The Asiatic subspecies is now found only in Iran in critically endangered numbers.The Asiatic cheetah whose long history on the Indian subcontinent gave the Sanskrit-derived vernacular name \"cheetah\", or \"spotted\", to the entire species,  Acinonyx jubatus, also had a gradual history of habitat loss there. In Punjab, before the thorn forests were cleared for agriculture and human settlement,  they were intermixed with open grasslands grazed by large herds of blackbuck; these co-existed with their main natural predator, the Asiatic cheetah.  The blackbuck is no longer extant in Punjab. Later, more habitat loss, prey depletion, and trophy hunting were to lead to the extinction of the Asiatic cheetah in other regions of India.\\nDiscussions on cheetah reintroduction in India began soon after extinction was confirmed, in the mid-1950s. Proposals were made to the governments of Iran from the 1970s, but fell through chiefly for reasons of political instability there. Offers from Kenya for introducing African cheetahs were made as early as the 1980s. Proposals for the introduction of African cheetahs were made by the Indian government in 2009, but disallowed by India\\'s supreme court. The court reversed its decision in early 2020, allowing the import of a small number, on an experimental basis for testing long-term adaptation. On 17 September 2022, five female and three male southeast African cheetahs, between the ages of four and six (a gift from the government of Namibia), were released in a small quarantined enclosure within the Kuno National Park in the state of Madhya Pradesh. The cheetahs, all fitted with radio collars, will remain in the quarantined enclosure for a month; initially, the males (and later the females) will be released into the 748.76 km2 (289.10 sq mi) park. The relocation has been supervised by Yadvendradev V. Jhala of the Wildlife Institute of India and zoologist Laurie Marker, of the Namibia-based Cheetah Conservation Fund. Subsequently, 12 cheetahs from South Africa will be released in Kuno;   eventually, the total number of African cheetahs in Kuno will be brought up to 40 individuals.  As of Jan 16, 2024, seven adult cheetahs from Africa and three cubs (of four born in Kuno two months earlier) had died in Kuno National Park.\\nThe scientific reaction to the translocation has been mixed. Adrian Tordiffe (a wildlife veterinary pharmacologist at the University of Pretoria who will be supervising the release of the cheetahs) is an enthusiast,  who views India as providing \"protected space\" for the fragmented and threatened population of the world\\'s cheetahs.  K. Ullas Karanth, one of India\\'s tiger experts, has been critical of the effort, considering it to be a \"PR exercise.\" India\\'s \"realities\", he says, such as human overpopulation, and the presence of larger feline predators and packs of feral dogs, could all cause potentially \"high mortalities,\" and require a continual import of African cheetahs. Kuno National Park is a relatively new national park, having received that status in 2018. It had been founded previously as a wildlife sanctuary to implement the Asiatic Lion Reintroduction Project, which aimed to establish a second Asiatic lion population in India. The goal was to protect the isolated lions of the Gir National Park (in Gujarat) from a potential mass mortality event, set off by the outbreak of an epizootic. Although the state government of Gujarat was ordered by India\\'s Supreme Court in April 2013 to transfer a small population of lions from Gujarat to Kuno, and was given six months to complete the transfer, they ultimately  resisted implementing the order.', 'source': 'https://en.wikipedia.org/wiki/Cheetah_reintroduction_in_India'}),\n",
       "  Document(page_content='The Southeast African cheetah (Acinonyx jubatus jubatus) is the nominate cheetah subspecies native to East and Southern Africa. The Southern African cheetah lives mainly in the lowland areas and deserts of the Kalahari, the savannahs of Okavango Delta, and the grasslands of the Transvaal region in South Africa. In Namibia, cheetahs are mostly found in farmlands. In India, four cheetahs of the subspecies are living in Kuno National Park in Madhya Pradesh after having been introduced there.\\n\\n\\n== Taxonomy ==\\nThe Southern African cheetah was first described by German naturalist Johann Christian Daniel von Schreber in his book Die Säugethiere in Abbildungen nach der Natur mit Beschreibungen (The Mammals illustrated as in Nature with Descriptions), published in 1775. Schreber described the species on basis of a specimen from the Cape of Good Hope. It is therefore the nominate subspecies. Subpopulations have been called \"South African cheetah\" and \"Namibian cheetah.\"Following Schreber\\'s description, other naturalists and zoologists also described cheetah specimens from many parts of Southern and East Africa that today are all considered synonyms of A. j. jubatus:\\nFelis guttata proposed in 1804 by Johann Hermann;\\nFelis fearonii proposed in 1834 by Andrew Smith;\\nFelis lanea proposed in 1877 by Philip Sclater;\\nAcinonyx jubatus obergi proposed in 1913 by Max Hilzheimer;\\nAcinonyx jubatus ngorongorensis proposed in 1913 by Hilzheimer on basis of a specimen from Ngorongoro, German East Africa;\\nAcinonyx jubatus velox proposed in 1913 by Edmund Heller on basis of a cheetah that was shot by Kermit Roosevelt in June 1909 in the Kenyan highlands.\\nAcinonyx rex proposed in 1927 by Reginald Innes Pocock on basis of a specimen from the Umvukwe Range in Rhodesia.In 2005, the authors of Mammal Species of the World grouped A. j. guttata, A. j. lanea, A. j. obergi, and A. j. rex under A j. jubatus, whilst recognizing A. j. raineyi and A. j. velox as valid taxa and considering P. l. ngorongore', metadata={'title': 'Southeast African cheetah', 'summary': 'The Southeast African cheetah (Acinonyx jubatus jubatus) is the nominate cheetah subspecies native to East and Southern Africa. The Southern African cheetah lives mainly in the lowland areas and deserts of the Kalahari, the savannahs of Okavango Delta, and the grasslands of the Transvaal region in South Africa. In Namibia, cheetahs are mostly found in farmlands. In India, four cheetahs of the subspecies are living in Kuno National Park in Madhya Pradesh after having been introduced there.', 'source': 'https://en.wikipedia.org/wiki/Southeast_African_cheetah'}),\n",
       "  Document(page_content='Footspeed, or sprint speed, is the maximum speed at which a human can run. It is affected by many factors, varies greatly throughout the population, and is important in athletics and many sports, such as association football, rugby football, American football,  track and field, field hockey, tennis, baseball, and basketball.\\n\\n\\n== Factors in speed ==\\nThe key determinant of footspeed in sprinting is the predominance of one distinct type of muscle fibre over another, specifically the ratio of fast-twitch muscles to slow-twitch muscles in a sprinter\\'s physical makeup.  Though fast-twitch muscles produce no more energy than slow-twitch muscles when they contract, they do so more rapidly through a process of anaerobic metabolism, though at the cost of inferior efficiency over longer periods of firing.  The average human has an almost-equal ratio of fast-twitch to slow-twitch fibers, but top sprinters may have as much as 80% fast-twitch fibers, while top long-distance runners may have only 20%. This ratio is believed to have genetic origins, though some assert that it can be adjusted by muscle training.  \"Speed camps\" and \"Speed Training Manuals\", which purport to provide fractional increases in maximum footspeed, are popular among budding professional athletes, and some sources estimate that 17–19% of speed can be trained.Though good running form is useful in increasing speed, fast and slow runners have been shown to move their legs at nearly the same rate – it is the force exerted by the leg on the ground that separates fast sprinters from slow.  Top short-distance runners exert as much as four times their body weight in pressure on the running surface. For this reason, muscle mass in the legs, relative to total body weight, is a key factor in maximizing footspeed.\\n\\n\\n== Limits of speed ==\\nThe record is 44.72 km/h (27.78 mph), measured between meter 60 and meter 80 of the 100 meters sprint at the 2009 World Championships in Athletics by Usain Bolt. (Bolt\\'s average speed o', metadata={'title': 'Footspeed', 'summary': 'Footspeed, or sprint speed, is the maximum speed at which a human can run. It is affected by many factors, varies greatly throughout the population, and is important in athletics and many sports, such as association football, rugby football, American football,  track and field, field hockey, tennis, baseball, and basketball.\\n\\n', 'source': 'https://en.wikipedia.org/wiki/Footspeed'}),\n",
       "  Document(page_content=\"This is a list of the fastest animals in the world, by types of animal.\\n\\n\\n== Fastest organism ==\\nThe peregrine falcon is the fastest bird, and the fastest member of the animal kingdom, with a diving speed of over 300 km/h (190 mph). The fastest land animal is the cheetah. Among the fastest animals in the sea is the black marlin, with uncertain and conflicting reports of recorded speeds.When drawing comparisons between different classes of animals, an alternative unit is sometimes used for organisms: body length per second. On this basis the 'fastest' organism on earth, relative to its body length, is the Southern Californian mite, Paratarsotomus macropalpis, which has a speed of 322 body lengths per second. The equivalent speed for a human, running as fast as this mite, would be 1,300 mph (2,092 km/h), or approximately Mach 1.7. The speed of the P. macropalpis is far in excess of the previous record holder, the Australian tiger beetle Cicindela eburneola, which is the fastest insect in the world relative to body size, with a recorded speed of 1.86 metres per second (6.7 km/h; 4.2 mph), or 171 body lengths per second. The cheetah, the fastest land mammal, scores at only 16 body lengths per second, while Anna's hummingbird has the highest known length-specific velocity attained by any vertebrate.\\n\\n\\n== Invertebrates ==\\n\\n\\n== Fish ==\\nDue to physical constraints, fish may be incapable of exceeding swim speeds of 36 km/h (22 mph). The larger reported figures below are therefore highly questionable:\\n\\n\\n== Amphibians ==\\n\\n\\n== Reptiles ==\\n\\n\\n== Birds ==\\n\\n\\n== Mammals ==\\n\\n\\n== See also ==\\nSpeed records\\n\\n\\n== Notes ==\\n\\n\\n== References ==\", metadata={'title': 'Fastest animals', 'summary': 'This is a list of the fastest animals in the world, by types of animal.', 'source': 'https://en.wikipedia.org/wiki/Fastest_animals'}),\n",
       "  Document(page_content=\"Pursuit predation is a form of predation in which predators actively give chase to their prey, either solitarily or as a group.  It is an alternate predation strategy to ambush predation — pursuit predators rely on superior speed, endurance and/or teamwork to seize the prey, while ambush predators use concealment, luring, exploiting of surroundings and the element of surprise to capture the prey.  While the two patterns of predation are not mutually exclusive, morphological differences in an organism's body plan can create an evolutionary bias favoring either type of predation.\\nPursuit predation is typically observed in carnivorous species within the kingdom Animalia, such as cheetahs, lions, wolves and early Homo species.  The chase can be initiated either by the predator, or by the prey if it is alerted to a predator's presence and attempt to flee before the predator gets close.  The chase ends either when the predator successfully catches up and tackles the prey, or when the predator abandons the attempt after the prey outruns it and escapes.\\nOne particular form of pursuit predation is persistence hunting, where the predator stalks the prey slowly but persistently to wear it down physically with fatigue or overheating; some animals are examples of both types of pursuit.\\n\\n\\n== Strategy ==\\nThere is still uncertainty as to whether predators behave with a general tactic or strategy while preying. However, among pursuit predators there are several common behaviors. Often, predators will scout potential prey, assessing prey quantity and density prior to engaging in a pursuit. Certain predators choose to pursue prey primarily in a group of conspecifics; these animals are known as pack hunters or group pursuers. Other species choose to hunt alone. These two behaviors are typically due to differences in hunting success, where some groups are very successful in groups and others are more successful alone. Pursuit predators may also choose to either exhaust their metabolic r\", metadata={'title': 'Pursuit predation', 'summary': \"Pursuit predation is a form of predation in which predators actively give chase to their prey, either solitarily or as a group.  It is an alternate predation strategy to ambush predation — pursuit predators rely on superior speed, endurance and/or teamwork to seize the prey, while ambush predators use concealment, luring, exploiting of surroundings and the element of surprise to capture the prey.  While the two patterns of predation are not mutually exclusive, morphological differences in an organism's body plan can create an evolutionary bias favoring either type of predation.\\nPursuit predation is typically observed in carnivorous species within the kingdom Animalia, such as cheetahs, lions, wolves and early Homo species.  The chase can be initiated either by the predator, or by the prey if it is alerted to a predator's presence and attempt to flee before the predator gets close.  The chase ends either when the predator successfully catches up and tackles the prey, or when the predator abandons the attempt after the prey outruns it and escapes.\\nOne particular form of pursuit predation is persistence hunting, where the predator stalks the prey slowly but persistently to wear it down physically with fatigue or overheating; some animals are examples of both types of pursuit.\", 'source': 'https://en.wikipedia.org/wiki/Pursuit_predation'})]}"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain_1.invoke(\"How fast are cheetahs?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94f2898a-ef4d-423a-b002-910fef7a65c9",
   "metadata": {},
   "source": [
    "LangSmith trace: https://smith.langchain.com/public/e38081da-774b-493b-b193-dab7711f99e9/r"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdbd1407-8a5b-4c35-aa2b-9d26424edb93",
   "metadata": {},
   "source": [
    "### Cite snippets\n",
    "\n",
    "What if we want to cite actual text spans? We can try to get our model to return these, too.\n",
    "\n",
    "*Aside: Note that if we break up our documents so that we have many documents with only a sentence or two instead of a few long documents, citing documents becomes roughly equivalent to citing snippets, and may be easier for the model because the model just needs to return an identifier for each snippet instead of the actual text. Probably worth trying both approaches and evaluating.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "fbf708aa-e8ac-4dea-bb57-82229597e2e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Citation(BaseModel):\n",
    "    source_id: int = Field(\n",
    "        ...,\n",
    "        description=\"The integer ID of a SPECIFIC source which justifies the answer.\",\n",
    "    )\n",
    "    quote: str = Field(\n",
    "        ...,\n",
    "        description=\"The VERBATIM quote from the specified source that justifies the answer.\",\n",
    "    )\n",
    "\n",
    "\n",
    "class quoted_answer(BaseModel):\n",
    "    \"\"\"Answer the user question based only on the given sources, and cite the sources used.\"\"\"\n",
    "\n",
    "    answer: str = Field(\n",
    "        ...,\n",
    "        description=\"The answer to the user question, which is based only on the given sources.\",\n",
    "    )\n",
    "    citations: List[Citation] = Field(\n",
    "        ..., description=\"Citations from the given sources that justify the answer.\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a488bd73-2676-476a-a338-b69e5f5479da",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_parser_2 = JsonOutputKeyToolsParser(key_name=\"quoted_answer\", return_single=True)\n",
    "llm_with_tool_2 = llm.bind_tools(\n",
    "    [quoted_answer],\n",
    "    tool_choice=\"quoted_answer\",\n",
    ")\n",
    "format_2 = itemgetter(\"docs\") | RunnableLambda(format_docs_with_id)\n",
    "answer_2 = prompt | llm_with_tool_2 | output_parser_2\n",
    "chain_2 = (\n",
    "    RunnableParallel(question=RunnablePassthrough(), docs=wiki)\n",
    "    .assign(context=format_2)\n",
    "    .assign(quoted_answer=answer_2)\n",
    "    .pick([\"quoted_answer\", \"docs\"])\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "232fb234-2444-4c3d-9fbb-1893dec2e8d0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'quoted_answer': {'answer': 'Cheetahs can run at speeds of 93 to 104 km/h (58 to 65 mph).',\n",
       "  'citations': [{'source_id': 0,\n",
       "    'quote': 'The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail.'}]},\n",
       " 'docs': [Document(page_content='The cheetah (Acinonyx jubatus) is a large cat and the fastest land animal. It has a tawny to creamy white or pale buff fur that is marked with evenly spaced, solid black spots. The head is small and rounded, with a short snout and black tear-like facial streaks. It reaches 67–94 cm (26–37 in) at the shoulder, and the head-and-body length is between 1.1 and 1.5 m (3 ft 7 in and 4 ft 11 in). Adults weigh between 21 and 72 kg (46 and 159 lb). The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail.\\nThe cheetah was first described in the late 18th century. Four subspecies are recognised today that are native to Africa and central Iran. An African subspecies was introduced to India in 2022. It is now distributed mainly in small, fragmented populations in northwestern, eastern and southern Africa and central Iran. It lives in a variety of habitats such as savannahs in the Serengeti, arid mountain ranges in the Sahara, and hilly desert terrain.\\nThe cheetah lives in three main social groups: females and their cubs, male \"coalitions\", and solitary males. While females lead a nomadic life searching for prey in large home ranges, males are more sedentary and instead establish much smaller territories in areas with plentiful prey and access to females. The cheetah is active during the day, with peaks during dawn and dusk. It feeds on small- to medium-sized prey, mostly weighing under 40 kg (88 lb), and prefers medium-sized ungulates such as impala, springbok and Thomson\\'s gazelles. The cheetah typically stalks its prey within 60–100 m (200–330 ft) before charging towards it, trips it during the chase and bites its throat to suffocate it to death. It breeds throughout the year. After a gestation of nearly three months, females give birth to a litter of three or four cubs. Cheetah cubs are highly vulnerable to predation by other large carnivores. They are weaned a', metadata={'title': 'Cheetah', 'summary': 'The cheetah (Acinonyx jubatus) is a large cat and the fastest land animal. It has a tawny to creamy white or pale buff fur that is marked with evenly spaced, solid black spots. The head is small and rounded, with a short snout and black tear-like facial streaks. It reaches 67–94 cm (26–37 in) at the shoulder, and the head-and-body length is between 1.1 and 1.5 m (3 ft 7 in and 4 ft 11 in). Adults weigh between 21 and 72 kg (46 and 159 lb). The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail.\\nThe cheetah was first described in the late 18th century. Four subspecies are recognised today that are native to Africa and central Iran. An African subspecies was introduced to India in 2022. It is now distributed mainly in small, fragmented populations in northwestern, eastern and southern Africa and central Iran. It lives in a variety of habitats such as savannahs in the Serengeti, arid mountain ranges in the Sahara, and hilly desert terrain.\\nThe cheetah lives in three main social groups: females and their cubs, male \"coalitions\", and solitary males. While females lead a nomadic life searching for prey in large home ranges, males are more sedentary and instead establish much smaller territories in areas with plentiful prey and access to females. The cheetah is active during the day, with peaks during dawn and dusk. It feeds on small- to medium-sized prey, mostly weighing under 40 kg (88 lb), and prefers medium-sized ungulates such as impala, springbok and Thomson\\'s gazelles. The cheetah typically stalks its prey within 60–100 m (200–330 ft) before charging towards it, trips it during the chase and bites its throat to suffocate it to death. It breeds throughout the year. After a gestation of nearly three months, females give birth to a litter of three or four cubs. Cheetah cubs are highly vulnerable to predation by other large carnivores. They are weaned at around four months and are independent by around 20 months of age.\\nThe cheetah is threatened by habitat loss, conflict with humans, poaching and high susceptibility to diseases. In 2016, the global cheetah population was estimated at 7,100 individuals in the wild; it is listed as Vulnerable on the IUCN Red List. It has been widely depicted in art, literature, advertising, and animation. It was tamed in ancient Egypt and trained for hunting ungulates in the Arabian Peninsula and India. It has been kept in zoos since the early 19th century.\\n\\n', 'source': 'https://en.wikipedia.org/wiki/Cheetah'}),\n",
       "  Document(page_content='More than 70 years after India\\'s native subspecies of the cheetah—the Asiatic cheetah (Acinonyx jubatus venaticus)—became extinct there, small numbers of Southeast African cheetah (Acinonyx jubatus jubatus) have been flown in from Namibia and South Africa to a national park in India.  The experiment has been permitted by India\\'s supreme court on a short-term basis to test long-term adaptation. The Asiatic subspecies is now found only in Iran in critically endangered numbers.The Asiatic cheetah whose long history on the Indian subcontinent gave the Sanskrit-derived vernacular name \"cheetah\", or \"spotted\", to the entire species,  Acinonyx jubatus, also had a gradual history of habitat loss there. In Punjab, before the thorn forests were cleared for agriculture and human settlement,  they were intermixed with open grasslands grazed by large herds of blackbuck; these co-existed with their main natural predator, the Asiatic cheetah.  The blackbuck is no longer extant in Punjab. Later, more habitat loss, prey depletion, and trophy hunting were to lead to the extinction of the Asiatic cheetah in other regions of India.\\nDiscussions on cheetah reintroduction in India began soon after extinction was confirmed, in the mid-1950s. Proposals were made to the governments of Iran from the 1970s, but fell through chiefly for reasons of political instability there. Offers from Kenya for introducing African cheetahs were made as early as the 1980s. Proposals for the introduction of African cheetahs were made by the Indian government in 2009, but disallowed by India\\'s supreme court. The court reversed its decision in early 2020, allowing the import of a small number, on an experimental basis for testing long-term adaptation. On 17 September 2022, five female and three male southeast African cheetahs, between the ages of four and six (a gift from the government of Namibia), were released in a small quarantined enclosure within the Kuno National Park in the state of Madhya Pradesh. The c', metadata={'title': 'Cheetah reintroduction in India', 'summary': 'More than 70 years after India\\'s native subspecies of the cheetah—the Asiatic cheetah (Acinonyx jubatus venaticus)—became extinct there, small numbers of Southeast African cheetah (Acinonyx jubatus jubatus) have been flown in from Namibia and South Africa to a national park in India.  The experiment has been permitted by India\\'s supreme court on a short-term basis to test long-term adaptation. The Asiatic subspecies is now found only in Iran in critically endangered numbers.The Asiatic cheetah whose long history on the Indian subcontinent gave the Sanskrit-derived vernacular name \"cheetah\", or \"spotted\", to the entire species,  Acinonyx jubatus, also had a gradual history of habitat loss there. In Punjab, before the thorn forests were cleared for agriculture and human settlement,  they were intermixed with open grasslands grazed by large herds of blackbuck; these co-existed with their main natural predator, the Asiatic cheetah.  The blackbuck is no longer extant in Punjab. Later, more habitat loss, prey depletion, and trophy hunting were to lead to the extinction of the Asiatic cheetah in other regions of India.\\nDiscussions on cheetah reintroduction in India began soon after extinction was confirmed, in the mid-1950s. Proposals were made to the governments of Iran from the 1970s, but fell through chiefly for reasons of political instability there. Offers from Kenya for introducing African cheetahs were made as early as the 1980s. Proposals for the introduction of African cheetahs were made by the Indian government in 2009, but disallowed by India\\'s supreme court. The court reversed its decision in early 2020, allowing the import of a small number, on an experimental basis for testing long-term adaptation. On 17 September 2022, five female and three male southeast African cheetahs, between the ages of four and six (a gift from the government of Namibia), were released in a small quarantined enclosure within the Kuno National Park in the state of Madhya Pradesh. The cheetahs, all fitted with radio collars, will remain in the quarantined enclosure for a month; initially, the males (and later the females) will be released into the 748.76 km2 (289.10 sq mi) park. The relocation has been supervised by Yadvendradev V. Jhala of the Wildlife Institute of India and zoologist Laurie Marker, of the Namibia-based Cheetah Conservation Fund. Subsequently, 12 cheetahs from South Africa will be released in Kuno;   eventually, the total number of African cheetahs in Kuno will be brought up to 40 individuals.  As of Jan 16, 2024, seven adult cheetahs from Africa and three cubs (of four born in Kuno two months earlier) had died in Kuno National Park.\\nThe scientific reaction to the translocation has been mixed. Adrian Tordiffe (a wildlife veterinary pharmacologist at the University of Pretoria who will be supervising the release of the cheetahs) is an enthusiast,  who views India as providing \"protected space\" for the fragmented and threatened population of the world\\'s cheetahs.  K. Ullas Karanth, one of India\\'s tiger experts, has been critical of the effort, considering it to be a \"PR exercise.\" India\\'s \"realities\", he says, such as human overpopulation, and the presence of larger feline predators and packs of feral dogs, could all cause potentially \"high mortalities,\" and require a continual import of African cheetahs. Kuno National Park is a relatively new national park, having received that status in 2018. It had been founded previously as a wildlife sanctuary to implement the Asiatic Lion Reintroduction Project, which aimed to establish a second Asiatic lion population in India. The goal was to protect the isolated lions of the Gir National Park (in Gujarat) from a potential mass mortality event, set off by the outbreak of an epizootic. Although the state government of Gujarat was ordered by India\\'s Supreme Court in April 2013 to transfer a small population of lions from Gujarat to Kuno, and was given six months to complete the transfer, they ultimately  resisted implementing the order.', 'source': 'https://en.wikipedia.org/wiki/Cheetah_reintroduction_in_India'}),\n",
       "  Document(page_content='The Southeast African cheetah (Acinonyx jubatus jubatus) is the nominate cheetah subspecies native to East and Southern Africa. The Southern African cheetah lives mainly in the lowland areas and deserts of the Kalahari, the savannahs of Okavango Delta, and the grasslands of the Transvaal region in South Africa. In Namibia, cheetahs are mostly found in farmlands. In India, four cheetahs of the subspecies are living in Kuno National Park in Madhya Pradesh after having been introduced there.\\n\\n\\n== Taxonomy ==\\nThe Southern African cheetah was first described by German naturalist Johann Christian Daniel von Schreber in his book Die Säugethiere in Abbildungen nach der Natur mit Beschreibungen (The Mammals illustrated as in Nature with Descriptions), published in 1775. Schreber described the species on basis of a specimen from the Cape of Good Hope. It is therefore the nominate subspecies. Subpopulations have been called \"South African cheetah\" and \"Namibian cheetah.\"Following Schreber\\'s description, other naturalists and zoologists also described cheetah specimens from many parts of Southern and East Africa that today are all considered synonyms of A. j. jubatus:\\nFelis guttata proposed in 1804 by Johann Hermann;\\nFelis fearonii proposed in 1834 by Andrew Smith;\\nFelis lanea proposed in 1877 by Philip Sclater;\\nAcinonyx jubatus obergi proposed in 1913 by Max Hilzheimer;\\nAcinonyx jubatus ngorongorensis proposed in 1913 by Hilzheimer on basis of a specimen from Ngorongoro, German East Africa;\\nAcinonyx jubatus velox proposed in 1913 by Edmund Heller on basis of a cheetah that was shot by Kermit Roosevelt in June 1909 in the Kenyan highlands.\\nAcinonyx rex proposed in 1927 by Reginald Innes Pocock on basis of a specimen from the Umvukwe Range in Rhodesia.In 2005, the authors of Mammal Species of the World grouped A. j. guttata, A. j. lanea, A. j. obergi, and A. j. rex under A j. jubatus, whilst recognizing A. j. raineyi and A. j. velox as valid taxa and considering P. l. ngorongore', metadata={'title': 'Southeast African cheetah', 'summary': 'The Southeast African cheetah (Acinonyx jubatus jubatus) is the nominate cheetah subspecies native to East and Southern Africa. The Southern African cheetah lives mainly in the lowland areas and deserts of the Kalahari, the savannahs of Okavango Delta, and the grasslands of the Transvaal region in South Africa. In Namibia, cheetahs are mostly found in farmlands. In India, four cheetahs of the subspecies are living in Kuno National Park in Madhya Pradesh after having been introduced there.', 'source': 'https://en.wikipedia.org/wiki/Southeast_African_cheetah'}),\n",
       "  Document(page_content='Footspeed, or sprint speed, is the maximum speed at which a human can run. It is affected by many factors, varies greatly throughout the population, and is important in athletics and many sports, such as association football, rugby football, American football,  track and field, field hockey, tennis, baseball, and basketball.\\n\\n\\n== Factors in speed ==\\nThe key determinant of footspeed in sprinting is the predominance of one distinct type of muscle fibre over another, specifically the ratio of fast-twitch muscles to slow-twitch muscles in a sprinter\\'s physical makeup.  Though fast-twitch muscles produce no more energy than slow-twitch muscles when they contract, they do so more rapidly through a process of anaerobic metabolism, though at the cost of inferior efficiency over longer periods of firing.  The average human has an almost-equal ratio of fast-twitch to slow-twitch fibers, but top sprinters may have as much as 80% fast-twitch fibers, while top long-distance runners may have only 20%. This ratio is believed to have genetic origins, though some assert that it can be adjusted by muscle training.  \"Speed camps\" and \"Speed Training Manuals\", which purport to provide fractional increases in maximum footspeed, are popular among budding professional athletes, and some sources estimate that 17–19% of speed can be trained.Though good running form is useful in increasing speed, fast and slow runners have been shown to move their legs at nearly the same rate – it is the force exerted by the leg on the ground that separates fast sprinters from slow.  Top short-distance runners exert as much as four times their body weight in pressure on the running surface. For this reason, muscle mass in the legs, relative to total body weight, is a key factor in maximizing footspeed.\\n\\n\\n== Limits of speed ==\\nThe record is 44.72 km/h (27.78 mph), measured between meter 60 and meter 80 of the 100 meters sprint at the 2009 World Championships in Athletics by Usain Bolt. (Bolt\\'s average speed o', metadata={'title': 'Footspeed', 'summary': 'Footspeed, or sprint speed, is the maximum speed at which a human can run. It is affected by many factors, varies greatly throughout the population, and is important in athletics and many sports, such as association football, rugby football, American football,  track and field, field hockey, tennis, baseball, and basketball.\\n\\n', 'source': 'https://en.wikipedia.org/wiki/Footspeed'}),\n",
       "  Document(page_content=\"This is a list of the fastest animals in the world, by types of animal.\\n\\n\\n== Fastest organism ==\\nThe peregrine falcon is the fastest bird, and the fastest member of the animal kingdom, with a diving speed of over 300 km/h (190 mph). The fastest land animal is the cheetah. Among the fastest animals in the sea is the black marlin, with uncertain and conflicting reports of recorded speeds.When drawing comparisons between different classes of animals, an alternative unit is sometimes used for organisms: body length per second. On this basis the 'fastest' organism on earth, relative to its body length, is the Southern Californian mite, Paratarsotomus macropalpis, which has a speed of 322 body lengths per second. The equivalent speed for a human, running as fast as this mite, would be 1,300 mph (2,092 km/h), or approximately Mach 1.7. The speed of the P. macropalpis is far in excess of the previous record holder, the Australian tiger beetle Cicindela eburneola, which is the fastest insect in the world relative to body size, with a recorded speed of 1.86 metres per second (6.7 km/h; 4.2 mph), or 171 body lengths per second. The cheetah, the fastest land mammal, scores at only 16 body lengths per second, while Anna's hummingbird has the highest known length-specific velocity attained by any vertebrate.\\n\\n\\n== Invertebrates ==\\n\\n\\n== Fish ==\\nDue to physical constraints, fish may be incapable of exceeding swim speeds of 36 km/h (22 mph). The larger reported figures below are therefore highly questionable:\\n\\n\\n== Amphibians ==\\n\\n\\n== Reptiles ==\\n\\n\\n== Birds ==\\n\\n\\n== Mammals ==\\n\\n\\n== See also ==\\nSpeed records\\n\\n\\n== Notes ==\\n\\n\\n== References ==\", metadata={'title': 'Fastest animals', 'summary': 'This is a list of the fastest animals in the world, by types of animal.', 'source': 'https://en.wikipedia.org/wiki/Fastest_animals'}),\n",
       "  Document(page_content=\"Pursuit predation is a form of predation in which predators actively give chase to their prey, either solitarily or as a group.  It is an alternate predation strategy to ambush predation — pursuit predators rely on superior speed, endurance and/or teamwork to seize the prey, while ambush predators use concealment, luring, exploiting of surroundings and the element of surprise to capture the prey.  While the two patterns of predation are not mutually exclusive, morphological differences in an organism's body plan can create an evolutionary bias favoring either type of predation.\\nPursuit predation is typically observed in carnivorous species within the kingdom Animalia, such as cheetahs, lions, wolves and early Homo species.  The chase can be initiated either by the predator, or by the prey if it is alerted to a predator's presence and attempt to flee before the predator gets close.  The chase ends either when the predator successfully catches up and tackles the prey, or when the predator abandons the attempt after the prey outruns it and escapes.\\nOne particular form of pursuit predation is persistence hunting, where the predator stalks the prey slowly but persistently to wear it down physically with fatigue or overheating; some animals are examples of both types of pursuit.\\n\\n\\n== Strategy ==\\nThere is still uncertainty as to whether predators behave with a general tactic or strategy while preying. However, among pursuit predators there are several common behaviors. Often, predators will scout potential prey, assessing prey quantity and density prior to engaging in a pursuit. Certain predators choose to pursue prey primarily in a group of conspecifics; these animals are known as pack hunters or group pursuers. Other species choose to hunt alone. These two behaviors are typically due to differences in hunting success, where some groups are very successful in groups and others are more successful alone. Pursuit predators may also choose to either exhaust their metabolic r\", metadata={'title': 'Pursuit predation', 'summary': \"Pursuit predation is a form of predation in which predators actively give chase to their prey, either solitarily or as a group.  It is an alternate predation strategy to ambush predation — pursuit predators rely on superior speed, endurance and/or teamwork to seize the prey, while ambush predators use concealment, luring, exploiting of surroundings and the element of surprise to capture the prey.  While the two patterns of predation are not mutually exclusive, morphological differences in an organism's body plan can create an evolutionary bias favoring either type of predation.\\nPursuit predation is typically observed in carnivorous species within the kingdom Animalia, such as cheetahs, lions, wolves and early Homo species.  The chase can be initiated either by the predator, or by the prey if it is alerted to a predator's presence and attempt to flee before the predator gets close.  The chase ends either when the predator successfully catches up and tackles the prey, or when the predator abandons the attempt after the prey outruns it and escapes.\\nOne particular form of pursuit predation is persistence hunting, where the predator stalks the prey slowly but persistently to wear it down physically with fatigue or overheating; some animals are examples of both types of pursuit.\", 'source': 'https://en.wikipedia.org/wiki/Pursuit_predation'})]}"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain_2.invoke(\"How fast are cheetahs?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28676cf1-4a2e-44d2-8b2f-36303a12a371",
   "metadata": {},
   "source": [
    "LangSmith trace: https://smith.langchain.com/public/ed19ea8d-5b99-4ebe-9809-9a3b4db6b39d/r"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb2d90a4-0370-4598-9f4b-e8e9a554346e",
   "metadata": {},
   "source": [
    "## Direct prompting\n",
    "\n",
    "Most models don't yet support function-calling. We can achieve similar results with direct prompting. Let's see what this looks like using an Anthropic chat model that is particularly proficient in working with XML:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "4e95bd8a-2f15-4e20-a1d9-225974b8d598",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_anthropic import ChatAnthropicMessages\n",
    "\n",
    "anthropic = ChatAnthropicMessages(model_name=\"claude-instant-1.2\")\n",
    "system = \"\"\"You're a helpful AI assistant. Given a user question and some Wikipedia article snippets, \\\n",
    "answer the user question and provide citations. If none of the articles answer the question, just say you don't know.\n",
    "\n",
    "Remember, you must return both an answer and citations. A citation consists of a VERBATIM quote that \\\n",
    "justifies the answer and the ID of the quote article. Return a citation for every quote across all articles \\\n",
    "that justify the answer. Use the following format for your final output:\n",
    "\n",
    "<cited_answer>\n",
    "    <answer></answer>\n",
    "    <citations>\n",
    "        <citation><source_id></source_id><quote></quote></citation>\n",
    "        <citation><source_id></source_id><quote></quote></citation>\n",
    "        ...\n",
    "    </citations>\n",
    "</cited_answer>\n",
    "\n",
    "Here are the Wikipedia articles:{context}\"\"\"\n",
    "prompt_3 = ChatPromptTemplate.from_messages(\n",
    "    [(\"system\", system), (\"human\", \"{question}\")]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "5861ca8c-63b7-4918-bdc6-fe4e53fe03ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.output_parsers import XMLOutputParser\n",
    "\n",
    "\n",
    "def format_docs_xml(docs: List[Document]) -> str:\n",
    "    formatted = []\n",
    "    for i, doc in enumerate(docs):\n",
    "        doc_str = f\"\"\"\\\n",
    "    <source id=\\\"{i}\\\">\n",
    "        <title>{doc.metadata['title']}</title>\n",
    "        <article_snippet>{doc.page_content}</article_snippet>\n",
    "    </source>\"\"\"\n",
    "        formatted.append(doc_str)\n",
    "    return \"\\n\\n<sources>\" + \"\\n\".join(formatted) + \"</sources>\"\n",
    "\n",
    "\n",
    "format_3 = itemgetter(\"docs\") | RunnableLambda(format_docs_xml)\n",
    "answer_3 = prompt_3 | anthropic | XMLOutputParser() | itemgetter(\"cited_answer\")\n",
    "chain_3 = (\n",
    "    RunnableParallel(question=RunnablePassthrough(), docs=wiki)\n",
    "    .assign(context=format_3)\n",
    "    .assign(cited_answer=answer_3)\n",
    "    .pick([\"cited_answer\", \"docs\"])\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "a3606806-929d-4b7b-b1bc-831db744a554",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'cited_answer': [{'answer': 'Cheetahs are the fastest land animals. They are capable of running at speeds of between 93 to 104 km/h (58 to 65 mph).'},\n",
       "  {'citations': [{'citation': [{'source_id': '0'},\n",
       "      {'quote': 'The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail.'}]}]}],\n",
       " 'docs': [Document(page_content='The cheetah (Acinonyx jubatus) is a large cat and the fastest land animal. It has a tawny to creamy white or pale buff fur that is marked with evenly spaced, solid black spots. The head is small and rounded, with a short snout and black tear-like facial streaks. It reaches 67–94 cm (26–37 in) at the shoulder, and the head-and-body length is between 1.1 and 1.5 m (3 ft 7 in and 4 ft 11 in). Adults weigh between 21 and 72 kg (46 and 159 lb). The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail.\\nThe cheetah was first described in the late 18th century. Four subspecies are recognised today that are native to Africa and central Iran. An African subspecies was introduced to India in 2022. It is now distributed mainly in small, fragmented populations in northwestern, eastern and southern Africa and central Iran. It lives in a variety of habitats such as savannahs in the Serengeti, arid mountain ranges in the Sahara, and hilly desert terrain.\\nThe cheetah lives in three main social groups: females and their cubs, male \"coalitions\", and solitary males. While females lead a nomadic life searching for prey in large home ranges, males are more sedentary and instead establish much smaller territories in areas with plentiful prey and access to females. The cheetah is active during the day, with peaks during dawn and dusk. It feeds on small- to medium-sized prey, mostly weighing under 40 kg (88 lb), and prefers medium-sized ungulates such as impala, springbok and Thomson\\'s gazelles. The cheetah typically stalks its prey within 60–100 m (200–330 ft) before charging towards it, trips it during the chase and bites its throat to suffocate it to death. It breeds throughout the year. After a gestation of nearly three months, females give birth to a litter of three or four cubs. Cheetah cubs are highly vulnerable to predation by other large carnivores. They are weaned a', metadata={'title': 'Cheetah', 'summary': 'The cheetah (Acinonyx jubatus) is a large cat and the fastest land animal. It has a tawny to creamy white or pale buff fur that is marked with evenly spaced, solid black spots. The head is small and rounded, with a short snout and black tear-like facial streaks. It reaches 67–94 cm (26–37 in) at the shoulder, and the head-and-body length is between 1.1 and 1.5 m (3 ft 7 in and 4 ft 11 in). Adults weigh between 21 and 72 kg (46 and 159 lb). The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail.\\nThe cheetah was first described in the late 18th century. Four subspecies are recognised today that are native to Africa and central Iran. An African subspecies was introduced to India in 2022. It is now distributed mainly in small, fragmented populations in northwestern, eastern and southern Africa and central Iran. It lives in a variety of habitats such as savannahs in the Serengeti, arid mountain ranges in the Sahara, and hilly desert terrain.\\nThe cheetah lives in three main social groups: females and their cubs, male \"coalitions\", and solitary males. While females lead a nomadic life searching for prey in large home ranges, males are more sedentary and instead establish much smaller territories in areas with plentiful prey and access to females. The cheetah is active during the day, with peaks during dawn and dusk. It feeds on small- to medium-sized prey, mostly weighing under 40 kg (88 lb), and prefers medium-sized ungulates such as impala, springbok and Thomson\\'s gazelles. The cheetah typically stalks its prey within 60–100 m (200–330 ft) before charging towards it, trips it during the chase and bites its throat to suffocate it to death. It breeds throughout the year. After a gestation of nearly three months, females give birth to a litter of three or four cubs. Cheetah cubs are highly vulnerable to predation by other large carnivores. They are weaned at around four months and are independent by around 20 months of age.\\nThe cheetah is threatened by habitat loss, conflict with humans, poaching and high susceptibility to diseases. In 2016, the global cheetah population was estimated at 7,100 individuals in the wild; it is listed as Vulnerable on the IUCN Red List. It has been widely depicted in art, literature, advertising, and animation. It was tamed in ancient Egypt and trained for hunting ungulates in the Arabian Peninsula and India. It has been kept in zoos since the early 19th century.\\n\\n', 'source': 'https://en.wikipedia.org/wiki/Cheetah'}),\n",
       "  Document(page_content='More than 70 years after India\\'s native subspecies of the cheetah—the Asiatic cheetah (Acinonyx jubatus venaticus)—became extinct there, small numbers of Southeast African cheetah (Acinonyx jubatus jubatus) have been flown in from Namibia and South Africa to a national park in India.  The experiment has been permitted by India\\'s supreme court on a short-term basis to test long-term adaptation. The Asiatic subspecies is now found only in Iran in critically endangered numbers.The Asiatic cheetah whose long history on the Indian subcontinent gave the Sanskrit-derived vernacular name \"cheetah\", or \"spotted\", to the entire species,  Acinonyx jubatus, also had a gradual history of habitat loss there. In Punjab, before the thorn forests were cleared for agriculture and human settlement,  they were intermixed with open grasslands grazed by large herds of blackbuck; these co-existed with their main natural predator, the Asiatic cheetah.  The blackbuck is no longer extant in Punjab. Later, more habitat loss, prey depletion, and trophy hunting were to lead to the extinction of the Asiatic cheetah in other regions of India.\\nDiscussions on cheetah reintroduction in India began soon after extinction was confirmed, in the mid-1950s. Proposals were made to the governments of Iran from the 1970s, but fell through chiefly for reasons of political instability there. Offers from Kenya for introducing African cheetahs were made as early as the 1980s. Proposals for the introduction of African cheetahs were made by the Indian government in 2009, but disallowed by India\\'s supreme court. The court reversed its decision in early 2020, allowing the import of a small number, on an experimental basis for testing long-term adaptation. On 17 September 2022, five female and three male southeast African cheetahs, between the ages of four and six (a gift from the government of Namibia), were released in a small quarantined enclosure within the Kuno National Park in the state of Madhya Pradesh. The c', metadata={'title': 'Cheetah reintroduction in India', 'summary': 'More than 70 years after India\\'s native subspecies of the cheetah—the Asiatic cheetah (Acinonyx jubatus venaticus)—became extinct there, small numbers of Southeast African cheetah (Acinonyx jubatus jubatus) have been flown in from Namibia and South Africa to a national park in India.  The experiment has been permitted by India\\'s supreme court on a short-term basis to test long-term adaptation. The Asiatic subspecies is now found only in Iran in critically endangered numbers.The Asiatic cheetah whose long history on the Indian subcontinent gave the Sanskrit-derived vernacular name \"cheetah\", or \"spotted\", to the entire species,  Acinonyx jubatus, also had a gradual history of habitat loss there. In Punjab, before the thorn forests were cleared for agriculture and human settlement,  they were intermixed with open grasslands grazed by large herds of blackbuck; these co-existed with their main natural predator, the Asiatic cheetah.  The blackbuck is no longer extant in Punjab. Later, more habitat loss, prey depletion, and trophy hunting were to lead to the extinction of the Asiatic cheetah in other regions of India.\\nDiscussions on cheetah reintroduction in India began soon after extinction was confirmed, in the mid-1950s. Proposals were made to the governments of Iran from the 1970s, but fell through chiefly for reasons of political instability there. Offers from Kenya for introducing African cheetahs were made as early as the 1980s. Proposals for the introduction of African cheetahs were made by the Indian government in 2009, but disallowed by India\\'s supreme court. The court reversed its decision in early 2020, allowing the import of a small number, on an experimental basis for testing long-term adaptation. On 17 September 2022, five female and three male southeast African cheetahs, between the ages of four and six (a gift from the government of Namibia), were released in a small quarantined enclosure within the Kuno National Park in the state of Madhya Pradesh. The cheetahs, all fitted with radio collars, will remain in the quarantined enclosure for a month; initially, the males (and later the females) will be released into the 748.76 km2 (289.10 sq mi) park. The relocation has been supervised by Yadvendradev V. Jhala of the Wildlife Institute of India and zoologist Laurie Marker, of the Namibia-based Cheetah Conservation Fund. Subsequently, 12 cheetahs from South Africa will be released in Kuno;   eventually, the total number of African cheetahs in Kuno will be brought up to 40 individuals.  As of Jan 16, 2024, seven adult cheetahs from Africa and three cubs (of four born in Kuno two months earlier) had died in Kuno National Park.\\nThe scientific reaction to the translocation has been mixed. Adrian Tordiffe (a wildlife veterinary pharmacologist at the University of Pretoria who will be supervising the release of the cheetahs) is an enthusiast,  who views India as providing \"protected space\" for the fragmented and threatened population of the world\\'s cheetahs.  K. Ullas Karanth, one of India\\'s tiger experts, has been critical of the effort, considering it to be a \"PR exercise.\" India\\'s \"realities\", he says, such as human overpopulation, and the presence of larger feline predators and packs of feral dogs, could all cause potentially \"high mortalities,\" and require a continual import of African cheetahs. Kuno National Park is a relatively new national park, having received that status in 2018. It had been founded previously as a wildlife sanctuary to implement the Asiatic Lion Reintroduction Project, which aimed to establish a second Asiatic lion population in India. The goal was to protect the isolated lions of the Gir National Park (in Gujarat) from a potential mass mortality event, set off by the outbreak of an epizootic. Although the state government of Gujarat was ordered by India\\'s Supreme Court in April 2013 to transfer a small population of lions from Gujarat to Kuno, and was given six months to complete the transfer, they ultimately  resisted implementing the order.', 'source': 'https://en.wikipedia.org/wiki/Cheetah_reintroduction_in_India'}),\n",
       "  Document(page_content='The Southeast African cheetah (Acinonyx jubatus jubatus) is the nominate cheetah subspecies native to East and Southern Africa. The Southern African cheetah lives mainly in the lowland areas and deserts of the Kalahari, the savannahs of Okavango Delta, and the grasslands of the Transvaal region in South Africa. In Namibia, cheetahs are mostly found in farmlands. In India, four cheetahs of the subspecies are living in Kuno National Park in Madhya Pradesh after having been introduced there.\\n\\n\\n== Taxonomy ==\\nThe Southern African cheetah was first described by German naturalist Johann Christian Daniel von Schreber in his book Die Säugethiere in Abbildungen nach der Natur mit Beschreibungen (The Mammals illustrated as in Nature with Descriptions), published in 1775. Schreber described the species on basis of a specimen from the Cape of Good Hope. It is therefore the nominate subspecies. Subpopulations have been called \"South African cheetah\" and \"Namibian cheetah.\"Following Schreber\\'s description, other naturalists and zoologists also described cheetah specimens from many parts of Southern and East Africa that today are all considered synonyms of A. j. jubatus:\\nFelis guttata proposed in 1804 by Johann Hermann;\\nFelis fearonii proposed in 1834 by Andrew Smith;\\nFelis lanea proposed in 1877 by Philip Sclater;\\nAcinonyx jubatus obergi proposed in 1913 by Max Hilzheimer;\\nAcinonyx jubatus ngorongorensis proposed in 1913 by Hilzheimer on basis of a specimen from Ngorongoro, German East Africa;\\nAcinonyx jubatus velox proposed in 1913 by Edmund Heller on basis of a cheetah that was shot by Kermit Roosevelt in June 1909 in the Kenyan highlands.\\nAcinonyx rex proposed in 1927 by Reginald Innes Pocock on basis of a specimen from the Umvukwe Range in Rhodesia.In 2005, the authors of Mammal Species of the World grouped A. j. guttata, A. j. lanea, A. j. obergi, and A. j. rex under A j. jubatus, whilst recognizing A. j. raineyi and A. j. velox as valid taxa and considering P. l. ngorongore', metadata={'title': 'Southeast African cheetah', 'summary': 'The Southeast African cheetah (Acinonyx jubatus jubatus) is the nominate cheetah subspecies native to East and Southern Africa. The Southern African cheetah lives mainly in the lowland areas and deserts of the Kalahari, the savannahs of Okavango Delta, and the grasslands of the Transvaal region in South Africa. In Namibia, cheetahs are mostly found in farmlands. In India, four cheetahs of the subspecies are living in Kuno National Park in Madhya Pradesh after having been introduced there.', 'source': 'https://en.wikipedia.org/wiki/Southeast_African_cheetah'}),\n",
       "  Document(page_content='Footspeed, or sprint speed, is the maximum speed at which a human can run. It is affected by many factors, varies greatly throughout the population, and is important in athletics and many sports, such as association football, rugby football, American football,  track and field, field hockey, tennis, baseball, and basketball.\\n\\n\\n== Factors in speed ==\\nThe key determinant of footspeed in sprinting is the predominance of one distinct type of muscle fibre over another, specifically the ratio of fast-twitch muscles to slow-twitch muscles in a sprinter\\'s physical makeup.  Though fast-twitch muscles produce no more energy than slow-twitch muscles when they contract, they do so more rapidly through a process of anaerobic metabolism, though at the cost of inferior efficiency over longer periods of firing.  The average human has an almost-equal ratio of fast-twitch to slow-twitch fibers, but top sprinters may have as much as 80% fast-twitch fibers, while top long-distance runners may have only 20%. This ratio is believed to have genetic origins, though some assert that it can be adjusted by muscle training.  \"Speed camps\" and \"Speed Training Manuals\", which purport to provide fractional increases in maximum footspeed, are popular among budding professional athletes, and some sources estimate that 17–19% of speed can be trained.Though good running form is useful in increasing speed, fast and slow runners have been shown to move their legs at nearly the same rate – it is the force exerted by the leg on the ground that separates fast sprinters from slow.  Top short-distance runners exert as much as four times their body weight in pressure on the running surface. For this reason, muscle mass in the legs, relative to total body weight, is a key factor in maximizing footspeed.\\n\\n\\n== Limits of speed ==\\nThe record is 44.72 km/h (27.78 mph), measured between meter 60 and meter 80 of the 100 meters sprint at the 2009 World Championships in Athletics by Usain Bolt. (Bolt\\'s average speed o', metadata={'title': 'Footspeed', 'summary': 'Footspeed, or sprint speed, is the maximum speed at which a human can run. It is affected by many factors, varies greatly throughout the population, and is important in athletics and many sports, such as association football, rugby football, American football,  track and field, field hockey, tennis, baseball, and basketball.\\n\\n', 'source': 'https://en.wikipedia.org/wiki/Footspeed'}),\n",
       "  Document(page_content=\"This is a list of the fastest animals in the world, by types of animal.\\n\\n\\n== Fastest organism ==\\nThe peregrine falcon is the fastest bird, and the fastest member of the animal kingdom, with a diving speed of over 300 km/h (190 mph). The fastest land animal is the cheetah. Among the fastest animals in the sea is the black marlin, with uncertain and conflicting reports of recorded speeds.When drawing comparisons between different classes of animals, an alternative unit is sometimes used for organisms: body length per second. On this basis the 'fastest' organism on earth, relative to its body length, is the Southern Californian mite, Paratarsotomus macropalpis, which has a speed of 322 body lengths per second. The equivalent speed for a human, running as fast as this mite, would be 1,300 mph (2,092 km/h), or approximately Mach 1.7. The speed of the P. macropalpis is far in excess of the previous record holder, the Australian tiger beetle Cicindela eburneola, which is the fastest insect in the world relative to body size, with a recorded speed of 1.86 metres per second (6.7 km/h; 4.2 mph), or 171 body lengths per second. The cheetah, the fastest land mammal, scores at only 16 body lengths per second, while Anna's hummingbird has the highest known length-specific velocity attained by any vertebrate.\\n\\n\\n== Invertebrates ==\\n\\n\\n== Fish ==\\nDue to physical constraints, fish may be incapable of exceeding swim speeds of 36 km/h (22 mph). The larger reported figures below are therefore highly questionable:\\n\\n\\n== Amphibians ==\\n\\n\\n== Reptiles ==\\n\\n\\n== Birds ==\\n\\n\\n== Mammals ==\\n\\n\\n== See also ==\\nSpeed records\\n\\n\\n== Notes ==\\n\\n\\n== References ==\", metadata={'title': 'Fastest animals', 'summary': 'This is a list of the fastest animals in the world, by types of animal.', 'source': 'https://en.wikipedia.org/wiki/Fastest_animals'}),\n",
       "  Document(page_content=\"Pursuit predation is a form of predation in which predators actively give chase to their prey, either solitarily or as a group.  It is an alternate predation strategy to ambush predation — pursuit predators rely on superior speed, endurance and/or teamwork to seize the prey, while ambush predators use concealment, luring, exploiting of surroundings and the element of surprise to capture the prey.  While the two patterns of predation are not mutually exclusive, morphological differences in an organism's body plan can create an evolutionary bias favoring either type of predation.\\nPursuit predation is typically observed in carnivorous species within the kingdom Animalia, such as cheetahs, lions, wolves and early Homo species.  The chase can be initiated either by the predator, or by the prey if it is alerted to a predator's presence and attempt to flee before the predator gets close.  The chase ends either when the predator successfully catches up and tackles the prey, or when the predator abandons the attempt after the prey outruns it and escapes.\\nOne particular form of pursuit predation is persistence hunting, where the predator stalks the prey slowly but persistently to wear it down physically with fatigue or overheating; some animals are examples of both types of pursuit.\\n\\n\\n== Strategy ==\\nThere is still uncertainty as to whether predators behave with a general tactic or strategy while preying. However, among pursuit predators there are several common behaviors. Often, predators will scout potential prey, assessing prey quantity and density prior to engaging in a pursuit. Certain predators choose to pursue prey primarily in a group of conspecifics; these animals are known as pack hunters or group pursuers. Other species choose to hunt alone. These two behaviors are typically due to differences in hunting success, where some groups are very successful in groups and others are more successful alone. Pursuit predators may also choose to either exhaust their metabolic r\", metadata={'title': 'Pursuit predation', 'summary': \"Pursuit predation is a form of predation in which predators actively give chase to their prey, either solitarily or as a group.  It is an alternate predation strategy to ambush predation — pursuit predators rely on superior speed, endurance and/or teamwork to seize the prey, while ambush predators use concealment, luring, exploiting of surroundings and the element of surprise to capture the prey.  While the two patterns of predation are not mutually exclusive, morphological differences in an organism's body plan can create an evolutionary bias favoring either type of predation.\\nPursuit predation is typically observed in carnivorous species within the kingdom Animalia, such as cheetahs, lions, wolves and early Homo species.  The chase can be initiated either by the predator, or by the prey if it is alerted to a predator's presence and attempt to flee before the predator gets close.  The chase ends either when the predator successfully catches up and tackles the prey, or when the predator abandons the attempt after the prey outruns it and escapes.\\nOne particular form of pursuit predation is persistence hunting, where the predator stalks the prey slowly but persistently to wear it down physically with fatigue or overheating; some animals are examples of both types of pursuit.\", 'source': 'https://en.wikipedia.org/wiki/Pursuit_predation'})]}"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain_3.invoke(\"How fast are cheetahs?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "940db8d5-8f43-44dd-9738-04fc7464baac",
   "metadata": {},
   "source": [
    "LangSmith trace: https://smith.langchain.com/public/54bd9284-0a32-4a29-8540-ff72142f0d3d/r"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d4180b0-5d29-4bfa-85be-2a6161a872c4",
   "metadata": {},
   "source": [
    "## Retrieval post-processing\n",
    "\n",
    "Another approach is to post-process our retrieved documents to compress the content, so that the source content is already minimal enough that we don't need the model to cite specific sources or spans. For example, we could break up each document into a sentence or two, embed those and keep only the most relevant ones. LangChain has some built-in components for this. Here we'll use a [RecursiveCharacterTextSplitter](https://api.python.langchain.com/en/latest/text_splitter/langchain_text_splitters.RecursiveCharacterTextSplitter.html#langchain_text_splitters.RecursiveCharacterTextSplitter), which creates chunks of a sepacified size by splitting on separator substrings, and an [EmbeddingsFilter](https://api.python.langchain.com/en/latest/retrievers/langchain.retrievers.document_compressors.embeddings_filter.EmbeddingsFilter.html#langchain.retrievers.document_compressors.embeddings_filter.EmbeddingsFilter), which keeps only the texts with the most relevant embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "9b14f817-4454-47b2-9eb0-2b8783a8c252",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Adults weigh between 21 and 72 kg (46 and 159 lb). The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail\n",
      "\n",
      "\n",
      "\n",
      "The cheetah (Acinonyx jubatus) is a large cat and the fastest land animal. It has a tawny to creamy white or pale buff fur that is marked with evenly spaced, solid black spots. The head is small and rounded, with a short snout and black tear-like facial streaks. It reaches 67–94 cm (26–37 in) at the shoulder, and the head-and-body length is between 1.1 and 1.5 m (3 ft 7 in and 4 ft 11 in)\n",
      "\n",
      "\n",
      "\n",
      "2 mph), or 171 body lengths per second. The cheetah, the fastest land mammal, scores at only 16 body lengths per second, while Anna's hummingbird has the highest known length-specific velocity attained by any vertebrate\n",
      "\n",
      "\n",
      "\n",
      "It feeds on small- to medium-sized prey, mostly weighing under 40 kg (88 lb), and prefers medium-sized ungulates such as impala, springbok and Thomson's gazelles. The cheetah typically stalks its prey within 60–100 m (200–330 ft) before charging towards it, trips it during the chase and bites its throat to suffocate it to death. It breeds throughout the year\n",
      "\n",
      "\n",
      "\n",
      "The cheetah was first described in the late 18th century. Four subspecies are recognised today that are native to Africa and central Iran. An African subspecies was introduced to India in 2022. It is now distributed mainly in small, fragmented populations in northwestern, eastern and southern Africa and central Iran\n",
      "\n",
      "\n",
      "\n",
      "The cheetah lives in three main social groups: females and their cubs, male \"coalitions\", and solitary males. While females lead a nomadic life searching for prey in large home ranges, males are more sedentary and instead establish much smaller territories in areas with plentiful prey and access to females. The cheetah is active during the day, with peaks during dawn and dusk\n",
      "\n",
      "\n",
      "\n",
      "The peregrine falcon is the fastest bird, and the fastest member of the animal kingdom, with a diving speed of over 300 km/h (190 mph). The fastest land animal is the cheetah. Among the fastest animals in the sea is the black marlin, with uncertain and conflicting reports of recorded speeds\n",
      "\n",
      "\n",
      "\n",
      "Acinonyx jubatus velox proposed in 1913 by Edmund Heller on basis of a cheetah that was shot by Kermit Roosevelt in June 1909 in the Kenyan highlands.\n",
      "\n",
      "\n",
      "\n",
      "The Southeast African cheetah (Acinonyx jubatus jubatus) is the nominate cheetah subspecies native to East and Southern Africa. The Southern African cheetah lives mainly in the lowland areas and deserts of the Kalahari, the savannahs of Okavango Delta, and the grasslands of the Transvaal region in South Africa. In Namibia, cheetahs are mostly found in farmlands\n",
      "\n",
      "\n",
      "\n",
      "On 17 September 2022, five female and three male southeast African cheetahs, between the ages of four and six (a gift from the government of Namibia), were released in a small quarantined enclosure within the Kuno National Park in the state of Madhya Pradesh. The c\n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from langchain.retrievers.document_compressors import EmbeddingsFilter\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "\n",
    "splitter = RecursiveCharacterTextSplitter(\n",
    "    chunk_size=400,\n",
    "    chunk_overlap=0,\n",
    "    separators=[\"\\n\\n\", \"\\n\", \".\", \" \"],\n",
    "    keep_separator=False,\n",
    ")\n",
    "compressor = EmbeddingsFilter(embeddings=OpenAIEmbeddings(), k=10)\n",
    "\n",
    "\n",
    "def split_and_filter(input) -> List[Document]:\n",
    "    docs = input[\"docs\"]\n",
    "    question = input[\"question\"]\n",
    "    split_docs = splitter.split_documents(docs)\n",
    "    stateful_docs = compressor.compress_documents(split_docs, question)\n",
    "    return [stateful_doc for stateful_doc in stateful_docs]\n",
    "\n",
    "\n",
    "retrieve = (\n",
    "    RunnableParallel(question=RunnablePassthrough(), docs=wiki) | split_and_filter\n",
    ")\n",
    "docs = retrieve.invoke(\"How fast are cheetahs?\")\n",
    "for doc in docs:\n",
    "    print(doc.page_content)\n",
    "    print(\"\\n\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "542fe2fc-06a3-4c31-9c02-a46daa6f8481",
   "metadata": {},
   "outputs": [],
   "source": [
    "chain_4 = (\n",
    "    RunnableParallel(question=RunnablePassthrough(), docs=retrieve)\n",
    "    .assign(context=format)\n",
    "    .assign(answer=answer)\n",
    "    .pick([\"answer\", \"docs\"])\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "id": "5bcb8cb2-daf3-48d4-b1d3-74b7066eb24c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'answer': 'Cheetahs are capable of running at speeds between 93 and 104 km/h (58 to 65 mph). They have evolved specialized adaptations for speed, including a light build, long thin legs, and a long tail.',\n",
       " 'docs': [Document(page_content='Adults weigh between 21 and 72 kg (46 and 159 lb). The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail', metadata={'title': 'Cheetah', 'summary': 'The cheetah (Acinonyx jubatus) is a large cat and the fastest land animal. It has a tawny to creamy white or pale buff fur that is marked with evenly spaced, solid black spots. The head is small and rounded, with a short snout and black tear-like facial streaks. It reaches 67–94 cm (26–37 in) at the shoulder, and the head-and-body length is between 1.1 and 1.5 m (3 ft 7 in and 4 ft 11 in). Adults weigh between 21 and 72 kg (46 and 159 lb). The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail.\\nThe cheetah was first described in the late 18th century. Four subspecies are recognised today that are native to Africa and central Iran. An African subspecies was introduced to India in 2022. It is now distributed mainly in small, fragmented populations in northwestern, eastern and southern Africa and central Iran. It lives in a variety of habitats such as savannahs in the Serengeti, arid mountain ranges in the Sahara, and hilly desert terrain.\\nThe cheetah lives in three main social groups: females and their cubs, male \"coalitions\", and solitary males. While females lead a nomadic life searching for prey in large home ranges, males are more sedentary and instead establish much smaller territories in areas with plentiful prey and access to females. The cheetah is active during the day, with peaks during dawn and dusk. It feeds on small- to medium-sized prey, mostly weighing under 40 kg (88 lb), and prefers medium-sized ungulates such as impala, springbok and Thomson\\'s gazelles. The cheetah typically stalks its prey within 60–100 m (200–330 ft) before charging towards it, trips it during the chase and bites its throat to suffocate it to death. It breeds throughout the year. After a gestation of nearly three months, females give birth to a litter of three or four cubs. Cheetah cubs are highly vulnerable to predation by other large carnivores. They are weaned at around four months and are independent by around 20 months of age.\\nThe cheetah is threatened by habitat loss, conflict with humans, poaching and high susceptibility to diseases. In 2016, the global cheetah population was estimated at 7,100 individuals in the wild; it is listed as Vulnerable on the IUCN Red List. It has been widely depicted in art, literature, advertising, and animation. It was tamed in ancient Egypt and trained for hunting ungulates in the Arabian Peninsula and India. It has been kept in zoos since the early 19th century.', 'source': 'https://en.wikipedia.org/wiki/Cheetah'}),\n",
       "  Document(page_content='The cheetah (Acinonyx jubatus) is a large cat and the fastest land animal. It has a tawny to creamy white or pale buff fur that is marked with evenly spaced, solid black spots. The head is small and rounded, with a short snout and black tear-like facial streaks. It reaches 67–94 cm (26–37 in) at the shoulder, and the head-and-body length is between 1.1 and 1.5 m (3 ft 7 in and 4 ft 11 in)', metadata={'title': 'Cheetah', 'summary': 'The cheetah (Acinonyx jubatus) is a large cat and the fastest land animal. It has a tawny to creamy white or pale buff fur that is marked with evenly spaced, solid black spots. The head is small and rounded, with a short snout and black tear-like facial streaks. It reaches 67–94 cm (26–37 in) at the shoulder, and the head-and-body length is between 1.1 and 1.5 m (3 ft 7 in and 4 ft 11 in). Adults weigh between 21 and 72 kg (46 and 159 lb). The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail.\\nThe cheetah was first described in the late 18th century. Four subspecies are recognised today that are native to Africa and central Iran. An African subspecies was introduced to India in 2022. It is now distributed mainly in small, fragmented populations in northwestern, eastern and southern Africa and central Iran. It lives in a variety of habitats such as savannahs in the Serengeti, arid mountain ranges in the Sahara, and hilly desert terrain.\\nThe cheetah lives in three main social groups: females and their cubs, male \"coalitions\", and solitary males. While females lead a nomadic life searching for prey in large home ranges, males are more sedentary and instead establish much smaller territories in areas with plentiful prey and access to females. The cheetah is active during the day, with peaks during dawn and dusk. It feeds on small- to medium-sized prey, mostly weighing under 40 kg (88 lb), and prefers medium-sized ungulates such as impala, springbok and Thomson\\'s gazelles. The cheetah typically stalks its prey within 60–100 m (200–330 ft) before charging towards it, trips it during the chase and bites its throat to suffocate it to death. It breeds throughout the year. After a gestation of nearly three months, females give birth to a litter of three or four cubs. Cheetah cubs are highly vulnerable to predation by other large carnivores. They are weaned at around four months and are independent by around 20 months of age.\\nThe cheetah is threatened by habitat loss, conflict with humans, poaching and high susceptibility to diseases. In 2016, the global cheetah population was estimated at 7,100 individuals in the wild; it is listed as Vulnerable on the IUCN Red List. It has been widely depicted in art, literature, advertising, and animation. It was tamed in ancient Egypt and trained for hunting ungulates in the Arabian Peninsula and India. It has been kept in zoos since the early 19th century.', 'source': 'https://en.wikipedia.org/wiki/Cheetah'}),\n",
       "  Document(page_content=\"2 mph), or 171 body lengths per second. The cheetah, the fastest land mammal, scores at only 16 body lengths per second, while Anna's hummingbird has the highest known length-specific velocity attained by any vertebrate\", metadata={'title': 'Fastest animals', 'summary': 'This is a list of the fastest animals in the world, by types of animal.', 'source': 'https://en.wikipedia.org/wiki/Fastest_animals'}),\n",
       "  Document(page_content=\"It feeds on small- to medium-sized prey, mostly weighing under 40 kg (88 lb), and prefers medium-sized ungulates such as impala, springbok and Thomson's gazelles. The cheetah typically stalks its prey within 60–100 m (200–330 ft) before charging towards it, trips it during the chase and bites its throat to suffocate it to death. It breeds throughout the year\", metadata={'title': 'Cheetah', 'summary': 'The cheetah (Acinonyx jubatus) is a large cat and the fastest land animal. It has a tawny to creamy white or pale buff fur that is marked with evenly spaced, solid black spots. The head is small and rounded, with a short snout and black tear-like facial streaks. It reaches 67–94 cm (26–37 in) at the shoulder, and the head-and-body length is between 1.1 and 1.5 m (3 ft 7 in and 4 ft 11 in). Adults weigh between 21 and 72 kg (46 and 159 lb). The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail.\\nThe cheetah was first described in the late 18th century. Four subspecies are recognised today that are native to Africa and central Iran. An African subspecies was introduced to India in 2022. It is now distributed mainly in small, fragmented populations in northwestern, eastern and southern Africa and central Iran. It lives in a variety of habitats such as savannahs in the Serengeti, arid mountain ranges in the Sahara, and hilly desert terrain.\\nThe cheetah lives in three main social groups: females and their cubs, male \"coalitions\", and solitary males. While females lead a nomadic life searching for prey in large home ranges, males are more sedentary and instead establish much smaller territories in areas with plentiful prey and access to females. The cheetah is active during the day, with peaks during dawn and dusk. It feeds on small- to medium-sized prey, mostly weighing under 40 kg (88 lb), and prefers medium-sized ungulates such as impala, springbok and Thomson\\'s gazelles. The cheetah typically stalks its prey within 60–100 m (200–330 ft) before charging towards it, trips it during the chase and bites its throat to suffocate it to death. It breeds throughout the year. After a gestation of nearly three months, females give birth to a litter of three or four cubs. Cheetah cubs are highly vulnerable to predation by other large carnivores. They are weaned at around four months and are independent by around 20 months of age.\\nThe cheetah is threatened by habitat loss, conflict with humans, poaching and high susceptibility to diseases. In 2016, the global cheetah population was estimated at 7,100 individuals in the wild; it is listed as Vulnerable on the IUCN Red List. It has been widely depicted in art, literature, advertising, and animation. It was tamed in ancient Egypt and trained for hunting ungulates in the Arabian Peninsula and India. It has been kept in zoos since the early 19th century.', 'source': 'https://en.wikipedia.org/wiki/Cheetah'}),\n",
       "  Document(page_content='The cheetah was first described in the late 18th century. Four subspecies are recognised today that are native to Africa and central Iran. An African subspecies was introduced to India in 2022. It is now distributed mainly in small, fragmented populations in northwestern, eastern and southern Africa and central Iran', metadata={'title': 'Cheetah', 'summary': 'The cheetah (Acinonyx jubatus) is a large cat and the fastest land animal. It has a tawny to creamy white or pale buff fur that is marked with evenly spaced, solid black spots. The head is small and rounded, with a short snout and black tear-like facial streaks. It reaches 67–94 cm (26–37 in) at the shoulder, and the head-and-body length is between 1.1 and 1.5 m (3 ft 7 in and 4 ft 11 in). Adults weigh between 21 and 72 kg (46 and 159 lb). The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail.\\nThe cheetah was first described in the late 18th century. Four subspecies are recognised today that are native to Africa and central Iran. An African subspecies was introduced to India in 2022. It is now distributed mainly in small, fragmented populations in northwestern, eastern and southern Africa and central Iran. It lives in a variety of habitats such as savannahs in the Serengeti, arid mountain ranges in the Sahara, and hilly desert terrain.\\nThe cheetah lives in three main social groups: females and their cubs, male \"coalitions\", and solitary males. While females lead a nomadic life searching for prey in large home ranges, males are more sedentary and instead establish much smaller territories in areas with plentiful prey and access to females. The cheetah is active during the day, with peaks during dawn and dusk. It feeds on small- to medium-sized prey, mostly weighing under 40 kg (88 lb), and prefers medium-sized ungulates such as impala, springbok and Thomson\\'s gazelles. The cheetah typically stalks its prey within 60–100 m (200–330 ft) before charging towards it, trips it during the chase and bites its throat to suffocate it to death. It breeds throughout the year. After a gestation of nearly three months, females give birth to a litter of three or four cubs. Cheetah cubs are highly vulnerable to predation by other large carnivores. They are weaned at around four months and are independent by around 20 months of age.\\nThe cheetah is threatened by habitat loss, conflict with humans, poaching and high susceptibility to diseases. In 2016, the global cheetah population was estimated at 7,100 individuals in the wild; it is listed as Vulnerable on the IUCN Red List. It has been widely depicted in art, literature, advertising, and animation. It was tamed in ancient Egypt and trained for hunting ungulates in the Arabian Peninsula and India. It has been kept in zoos since the early 19th century.', 'source': 'https://en.wikipedia.org/wiki/Cheetah'}),\n",
       "  Document(page_content='The cheetah lives in three main social groups: females and their cubs, male \"coalitions\", and solitary males. While females lead a nomadic life searching for prey in large home ranges, males are more sedentary and instead establish much smaller territories in areas with plentiful prey and access to females. The cheetah is active during the day, with peaks during dawn and dusk', metadata={'title': 'Cheetah', 'summary': 'The cheetah (Acinonyx jubatus) is a large cat and the fastest land animal. It has a tawny to creamy white or pale buff fur that is marked with evenly spaced, solid black spots. The head is small and rounded, with a short snout and black tear-like facial streaks. It reaches 67–94 cm (26–37 in) at the shoulder, and the head-and-body length is between 1.1 and 1.5 m (3 ft 7 in and 4 ft 11 in). Adults weigh between 21 and 72 kg (46 and 159 lb). The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail.\\nThe cheetah was first described in the late 18th century. Four subspecies are recognised today that are native to Africa and central Iran. An African subspecies was introduced to India in 2022. It is now distributed mainly in small, fragmented populations in northwestern, eastern and southern Africa and central Iran. It lives in a variety of habitats such as savannahs in the Serengeti, arid mountain ranges in the Sahara, and hilly desert terrain.\\nThe cheetah lives in three main social groups: females and their cubs, male \"coalitions\", and solitary males. While females lead a nomadic life searching for prey in large home ranges, males are more sedentary and instead establish much smaller territories in areas with plentiful prey and access to females. The cheetah is active during the day, with peaks during dawn and dusk. It feeds on small- to medium-sized prey, mostly weighing under 40 kg (88 lb), and prefers medium-sized ungulates such as impala, springbok and Thomson\\'s gazelles. The cheetah typically stalks its prey within 60–100 m (200–330 ft) before charging towards it, trips it during the chase and bites its throat to suffocate it to death. It breeds throughout the year. After a gestation of nearly three months, females give birth to a litter of three or four cubs. Cheetah cubs are highly vulnerable to predation by other large carnivores. They are weaned at around four months and are independent by around 20 months of age.\\nThe cheetah is threatened by habitat loss, conflict with humans, poaching and high susceptibility to diseases. In 2016, the global cheetah population was estimated at 7,100 individuals in the wild; it is listed as Vulnerable on the IUCN Red List. It has been widely depicted in art, literature, advertising, and animation. It was tamed in ancient Egypt and trained for hunting ungulates in the Arabian Peninsula and India. It has been kept in zoos since the early 19th century.', 'source': 'https://en.wikipedia.org/wiki/Cheetah'}),\n",
       "  Document(page_content='The peregrine falcon is the fastest bird, and the fastest member of the animal kingdom, with a diving speed of over 300 km/h (190 mph). The fastest land animal is the cheetah. Among the fastest animals in the sea is the black marlin, with uncertain and conflicting reports of recorded speeds', metadata={'title': 'Fastest animals', 'summary': 'This is a list of the fastest animals in the world, by types of animal.', 'source': 'https://en.wikipedia.org/wiki/Fastest_animals'}),\n",
       "  Document(page_content='Acinonyx jubatus velox proposed in 1913 by Edmund Heller on basis of a cheetah that was shot by Kermit Roosevelt in June 1909 in the Kenyan highlands.', metadata={'title': 'Southeast African cheetah', 'summary': 'The Southeast African cheetah (Acinonyx jubatus jubatus) is the nominate cheetah subspecies native to East and Southern Africa. The Southern African cheetah lives mainly in the lowland areas and deserts of the Kalahari, the savannahs of Okavango Delta, and the grasslands of the Transvaal region in South Africa. In Namibia, cheetahs are mostly found in farmlands. In India, four cheetahs of the subspecies are living in Kuno National Park in Madhya Pradesh after having been introduced there.', 'source': 'https://en.wikipedia.org/wiki/Southeast_African_cheetah'}),\n",
       "  Document(page_content='The Southeast African cheetah (Acinonyx jubatus jubatus) is the nominate cheetah subspecies native to East and Southern Africa. The Southern African cheetah lives mainly in the lowland areas and deserts of the Kalahari, the savannahs of Okavango Delta, and the grasslands of the Transvaal region in South Africa. In Namibia, cheetahs are mostly found in farmlands', metadata={'title': 'Southeast African cheetah', 'summary': 'The Southeast African cheetah (Acinonyx jubatus jubatus) is the nominate cheetah subspecies native to East and Southern Africa. The Southern African cheetah lives mainly in the lowland areas and deserts of the Kalahari, the savannahs of Okavango Delta, and the grasslands of the Transvaal region in South Africa. In Namibia, cheetahs are mostly found in farmlands. In India, four cheetahs of the subspecies are living in Kuno National Park in Madhya Pradesh after having been introduced there.', 'source': 'https://en.wikipedia.org/wiki/Southeast_African_cheetah'}),\n",
       "  Document(page_content='On 17 September 2022, five female and three male southeast African cheetahs, between the ages of four and six (a gift from the government of Namibia), were released in a small quarantined enclosure within the Kuno National Park in the state of Madhya Pradesh. The c', metadata={'title': 'Cheetah reintroduction in India', 'summary': 'More than 70 years after India\\'s native subspecies of the cheetah—the Asiatic cheetah (Acinonyx jubatus venaticus)—became extinct there, small numbers of Southeast African cheetah (Acinonyx jubatus jubatus) have been flown in from Namibia and South Africa to a national park in India.  The experiment has been permitted by India\\'s supreme court on a short-term basis to test long-term adaptation. The Asiatic subspecies is now found only in Iran in critically endangered numbers.The Asiatic cheetah whose long history on the Indian subcontinent gave the Sanskrit-derived vernacular name \"cheetah\", or \"spotted\", to the entire species,  Acinonyx jubatus, also had a gradual history of habitat loss there. In Punjab, before the thorn forests were cleared for agriculture and human settlement,  they were intermixed with open grasslands grazed by large herds of blackbuck; these co-existed with their main natural predator, the Asiatic cheetah.  The blackbuck is no longer extant in Punjab. Later, more habitat loss, prey depletion, and trophy hunting were to lead to the extinction of the Asiatic cheetah in other regions of India.\\nDiscussions on cheetah reintroduction in India began soon after extinction was confirmed, in the mid-1950s. Proposals were made to the governments of Iran from the 1970s, but fell through chiefly for reasons of political instability there. Offers from Kenya for introducing African cheetahs were made as early as the 1980s. Proposals for the introduction of African cheetahs were made by the Indian government in 2009, but disallowed by India\\'s supreme court. The court reversed its decision in early 2020, allowing the import of a small number, on an experimental basis for testing long-term adaptation. On 17 September 2022, five female and three male southeast African cheetahs, between the ages of four and six (a gift from the government of Namibia), were released in a small quarantined enclosure within the Kuno National Park in the state of Madhya Pradesh. The cheetahs, all fitted with radio collars, will remain in the quarantined enclosure for a month; initially, the males (and later the females) will be released into the 748.76 km2 (289.10 sq mi) park. The relocation has been supervised by Yadvendradev V. Jhala of the Wildlife Institute of India and zoologist Laurie Marker, of the Namibia-based Cheetah Conservation Fund. Subsequently, 12 cheetahs from South Africa will be released in Kuno;   eventually, the total number of African cheetahs in Kuno will be brought up to 40 individuals.  As of Jan 16, 2024, seven adult cheetahs from Africa and three cubs (of four born in Kuno two months earlier) had died in Kuno National Park.\\nThe scientific reaction to the translocation has been mixed. Adrian Tordiffe (a wildlife veterinary pharmacologist at the University of Pretoria who will be supervising the release of the cheetahs) is an enthusiast,  who views India as providing \"protected space\" for the fragmented and threatened population of the world\\'s cheetahs.  K. Ullas Karanth, one of India\\'s tiger experts, has been critical of the effort, considering it to be a \"PR exercise.\" India\\'s \"realities\", he says, such as human overpopulation, and the presence of larger feline predators and packs of feral dogs, could all cause potentially \"high mortalities,\" and require a continual import of African cheetahs. Kuno National Park is a relatively new national park, having received that status in 2018. It had been founded previously as a wildlife sanctuary to implement the Asiatic Lion Reintroduction Project, which aimed to establish a second Asiatic lion population in India. The goal was to protect the isolated lions of the Gir National Park (in Gujarat) from a potential mass mortality event, set off by the outbreak of an epizootic. Although the state government of Gujarat was ordered by India\\'s Supreme Court in April 2013 to transfer a small population of lions from Gujarat to Kuno, and was given six months to complete the transfer, they ultimately  resisted implementing the order.', 'source': 'https://en.wikipedia.org/wiki/Cheetah_reintroduction_in_India'})]}"
      ]
     },
     "execution_count": 78,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Note the documents have an article \"summary\" in the metadata that is now much longer than the\n",
    "# actual document page content. This summary isn't actually passed to the model.\n",
    "chain_4.invoke(\"How fast are cheetahs?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88ab8fd6-f6b4-4ba5-b022-f10cca983490",
   "metadata": {},
   "source": [
    "LangSmith trace: https://smith.langchain.com/public/f6a7ea78-05f3-47f2-a3dc-72747d1a9c64/r"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "445722dc-2ecb-45a4-9d4d-c172d0a2fa7d",
   "metadata": {},
   "source": [
    "## Generation post-processing\n",
    "\n",
    "Another approach is to post-process our model generation. In this example we'll first generate just an answer, and then we'll ask the model to annotate it's own answer with citations. The downside of this approach is of course that it is slower and more expensive, because two model calls need to be made.\n",
    "\n",
    "Let's apply this to our initial chain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "daff5cb9-7639-4d30-b6e7-d795736a2b58",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Citation(BaseModel):\n",
    "    source_id: int = Field(\n",
    "        ...,\n",
    "        description=\"The integer ID of a SPECIFIC source which justifies the answer.\",\n",
    "    )\n",
    "    quote: str = Field(\n",
    "        ...,\n",
    "        description=\"The VERBATIM quote from the specified source that justifies the answer.\",\n",
    "    )\n",
    "\n",
    "\n",
    "class annotated_answer(BaseModel):\n",
    "    \"\"\"Annotate the answer to the user question with quote citations that justify the answer.\"\"\"\n",
    "\n",
    "    citations: List[Citation] = Field(\n",
    "        ..., description=\"Citations from the given sources that justify the answer.\"\n",
    "    )\n",
    "\n",
    "\n",
    "llm_with_tools_5 = llm.bind_tools(\n",
    "    [annotated_answer],\n",
    "    tool_choice=\"annotated_answer\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "id": "6f505eb9-db02-4c49-add3-1e469844d7ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import MessagesPlaceholder\n",
    "\n",
    "prompt_5 = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            \"You're a helpful AI assistant. Given a user question and some Wikipedia article snippets, answer the user question. If none of the articles answer the question, just say you don't know.\\n\\nHere are the Wikipedia articles:{context}\",\n",
    "        ),\n",
    "        (\"human\", \"{question}\"),\n",
    "        MessagesPlaceholder(\"chat_history\", optional=True),\n",
    "    ]\n",
    ")\n",
    "answer_5 = prompt_5 | llm\n",
    "annotation_chain = (\n",
    "    prompt_5\n",
    "    | llm_with_tools_5\n",
    "    | JsonOutputKeyToolsParser(key_name=\"annotated_answer\", return_single=True)\n",
    "    | itemgetter(\"citations\")\n",
    ")\n",
    "\n",
    "chain_5 = (\n",
    "    RunnableParallel(question=RunnablePassthrough(), docs=wiki)\n",
    "    .assign(context=format)\n",
    "    .assign(ai_message=answer_5)\n",
    "    .assign(\n",
    "        chat_history=(lambda x: [x[\"ai_message\"]]),\n",
    "        answer=(lambda x: x[\"ai_message\"].content),\n",
    "    )\n",
    "    .assign(annotations=annotation_chain)\n",
    "    .pick([\"answer\", \"docs\", \"annotations\"])\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "id": "eb11c422-09b3-4d5a-87eb-3bad2e73cf6c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'answer': 'Cheetahs are capable of running at speeds between 93 to 104 km/h (58 to 65 mph). They have evolved specialized adaptations for speed, including a light build, long thin legs, and a long tail.',\n",
       " 'docs': [Document(page_content='The cheetah (Acinonyx jubatus) is a large cat and the fastest land animal. It has a tawny to creamy white or pale buff fur that is marked with evenly spaced, solid black spots. The head is small and rounded, with a short snout and black tear-like facial streaks. It reaches 67–94 cm (26–37 in) at the shoulder, and the head-and-body length is between 1.1 and 1.5 m (3 ft 7 in and 4 ft 11 in). Adults weigh between 21 and 72 kg (46 and 159 lb). The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail.\\nThe cheetah was first described in the late 18th century. Four subspecies are recognised today that are native to Africa and central Iran. An African subspecies was introduced to India in 2022. It is now distributed mainly in small, fragmented populations in northwestern, eastern and southern Africa and central Iran. It lives in a variety of habitats such as savannahs in the Serengeti, arid mountain ranges in the Sahara, and hilly desert terrain.\\nThe cheetah lives in three main social groups: females and their cubs, male \"coalitions\", and solitary males. While females lead a nomadic life searching for prey in large home ranges, males are more sedentary and instead establish much smaller territories in areas with plentiful prey and access to females. The cheetah is active during the day, with peaks during dawn and dusk. It feeds on small- to medium-sized prey, mostly weighing under 40 kg (88 lb), and prefers medium-sized ungulates such as impala, springbok and Thomson\\'s gazelles. The cheetah typically stalks its prey within 60–100 m (200–330 ft) before charging towards it, trips it during the chase and bites its throat to suffocate it to death. It breeds throughout the year. After a gestation of nearly three months, females give birth to a litter of three or four cubs. Cheetah cubs are highly vulnerable to predation by other large carnivores. They are weaned a', metadata={'title': 'Cheetah', 'summary': 'The cheetah (Acinonyx jubatus) is a large cat and the fastest land animal. It has a tawny to creamy white or pale buff fur that is marked with evenly spaced, solid black spots. The head is small and rounded, with a short snout and black tear-like facial streaks. It reaches 67–94 cm (26–37 in) at the shoulder, and the head-and-body length is between 1.1 and 1.5 m (3 ft 7 in and 4 ft 11 in). Adults weigh between 21 and 72 kg (46 and 159 lb). The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail.\\nThe cheetah was first described in the late 18th century. Four subspecies are recognised today that are native to Africa and central Iran. An African subspecies was introduced to India in 2022. It is now distributed mainly in small, fragmented populations in northwestern, eastern and southern Africa and central Iran. It lives in a variety of habitats such as savannahs in the Serengeti, arid mountain ranges in the Sahara, and hilly desert terrain.\\nThe cheetah lives in three main social groups: females and their cubs, male \"coalitions\", and solitary males. While females lead a nomadic life searching for prey in large home ranges, males are more sedentary and instead establish much smaller territories in areas with plentiful prey and access to females. The cheetah is active during the day, with peaks during dawn and dusk. It feeds on small- to medium-sized prey, mostly weighing under 40 kg (88 lb), and prefers medium-sized ungulates such as impala, springbok and Thomson\\'s gazelles. The cheetah typically stalks its prey within 60–100 m (200–330 ft) before charging towards it, trips it during the chase and bites its throat to suffocate it to death. It breeds throughout the year. After a gestation of nearly three months, females give birth to a litter of three or four cubs. Cheetah cubs are highly vulnerable to predation by other large carnivores. They are weaned at around four months and are independent by around 20 months of age.\\nThe cheetah is threatened by habitat loss, conflict with humans, poaching and high susceptibility to diseases. In 2016, the global cheetah population was estimated at 7,100 individuals in the wild; it is listed as Vulnerable on the IUCN Red List. It has been widely depicted in art, literature, advertising, and animation. It was tamed in ancient Egypt and trained for hunting ungulates in the Arabian Peninsula and India. It has been kept in zoos since the early 19th century.', 'source': 'https://en.wikipedia.org/wiki/Cheetah'}),\n",
       "  Document(page_content='More than 70 years after India\\'s native subspecies of the cheetah—the Asiatic cheetah (Acinonyx jubatus venaticus)—became extinct there, small numbers of Southeast African cheetah (Acinonyx jubatus jubatus) have been flown in from Namibia and South Africa to a national park in India.  The experiment has been permitted by India\\'s supreme court on a short-term basis to test long-term adaptation. The Asiatic subspecies is now found only in Iran in critically endangered numbers.The Asiatic cheetah whose long history on the Indian subcontinent gave the Sanskrit-derived vernacular name \"cheetah\", or \"spotted\", to the entire species,  Acinonyx jubatus, also had a gradual history of habitat loss there. In Punjab, before the thorn forests were cleared for agriculture and human settlement,  they were intermixed with open grasslands grazed by large herds of blackbuck; these co-existed with their main natural predator, the Asiatic cheetah.  The blackbuck is no longer extant in Punjab. Later, more habitat loss, prey depletion, and trophy hunting were to lead to the extinction of the Asiatic cheetah in other regions of India.\\nDiscussions on cheetah reintroduction in India began soon after extinction was confirmed, in the mid-1950s. Proposals were made to the governments of Iran from the 1970s, but fell through chiefly for reasons of political instability there. Offers from Kenya for introducing African cheetahs were made as early as the 1980s. Proposals for the introduction of African cheetahs were made by the Indian government in 2009, but disallowed by India\\'s supreme court. The court reversed its decision in early 2020, allowing the import of a small number, on an experimental basis for testing long-term adaptation. On 17 September 2022, five female and three male southeast African cheetahs, between the ages of four and six (a gift from the government of Namibia), were released in a small quarantined enclosure within the Kuno National Park in the state of Madhya Pradesh. The c', metadata={'title': 'Cheetah reintroduction in India', 'summary': 'More than 70 years after India\\'s native subspecies of the cheetah—the Asiatic cheetah (Acinonyx jubatus venaticus)—became extinct there, small numbers of Southeast African cheetah (Acinonyx jubatus jubatus) have been flown in from Namibia and South Africa to a national park in India.  The experiment has been permitted by India\\'s supreme court on a short-term basis to test long-term adaptation. The Asiatic subspecies is now found only in Iran in critically endangered numbers.The Asiatic cheetah whose long history on the Indian subcontinent gave the Sanskrit-derived vernacular name \"cheetah\", or \"spotted\", to the entire species,  Acinonyx jubatus, also had a gradual history of habitat loss there. In Punjab, before the thorn forests were cleared for agriculture and human settlement,  they were intermixed with open grasslands grazed by large herds of blackbuck; these co-existed with their main natural predator, the Asiatic cheetah.  The blackbuck is no longer extant in Punjab. Later, more habitat loss, prey depletion, and trophy hunting were to lead to the extinction of the Asiatic cheetah in other regions of India.\\nDiscussions on cheetah reintroduction in India began soon after extinction was confirmed, in the mid-1950s. Proposals were made to the governments of Iran from the 1970s, but fell through chiefly for reasons of political instability there. Offers from Kenya for introducing African cheetahs were made as early as the 1980s. Proposals for the introduction of African cheetahs were made by the Indian government in 2009, but disallowed by India\\'s supreme court. The court reversed its decision in early 2020, allowing the import of a small number, on an experimental basis for testing long-term adaptation. On 17 September 2022, five female and three male southeast African cheetahs, between the ages of four and six (a gift from the government of Namibia), were released in a small quarantined enclosure within the Kuno National Park in the state of Madhya Pradesh. The cheetahs, all fitted with radio collars, will remain in the quarantined enclosure for a month; initially, the males (and later the females) will be released into the 748.76 km2 (289.10 sq mi) park. The relocation has been supervised by Yadvendradev V. Jhala of the Wildlife Institute of India and zoologist Laurie Marker, of the Namibia-based Cheetah Conservation Fund. Subsequently, 12 cheetahs from South Africa will be released in Kuno;   eventually, the total number of African cheetahs in Kuno will be brought up to 40 individuals.  As of Jan 16, 2024, seven adult cheetahs from Africa and three cubs (of four born in Kuno two months earlier) had died in Kuno National Park.\\nThe scientific reaction to the translocation has been mixed. Adrian Tordiffe (a wildlife veterinary pharmacologist at the University of Pretoria who will be supervising the release of the cheetahs) is an enthusiast,  who views India as providing \"protected space\" for the fragmented and threatened population of the world\\'s cheetahs.  K. Ullas Karanth, one of India\\'s tiger experts, has been critical of the effort, considering it to be a \"PR exercise.\" India\\'s \"realities\", he says, such as human overpopulation, and the presence of larger feline predators and packs of feral dogs, could all cause potentially \"high mortalities,\" and require a continual import of African cheetahs. Kuno National Park is a relatively new national park, having received that status in 2018. It had been founded previously as a wildlife sanctuary to implement the Asiatic Lion Reintroduction Project, which aimed to establish a second Asiatic lion population in India. The goal was to protect the isolated lions of the Gir National Park (in Gujarat) from a potential mass mortality event, set off by the outbreak of an epizootic. Although the state government of Gujarat was ordered by India\\'s Supreme Court in April 2013 to transfer a small population of lions from Gujarat to Kuno, and was given six months to complete the transfer, they ultimately  resisted implementing the order.', 'source': 'https://en.wikipedia.org/wiki/Cheetah_reintroduction_in_India'}),\n",
       "  Document(page_content='The Southeast African cheetah (Acinonyx jubatus jubatus) is the nominate cheetah subspecies native to East and Southern Africa. The Southern African cheetah lives mainly in the lowland areas and deserts of the Kalahari, the savannahs of Okavango Delta, and the grasslands of the Transvaal region in South Africa. In Namibia, cheetahs are mostly found in farmlands. In India, four cheetahs of the subspecies are living in Kuno National Park in Madhya Pradesh after having been introduced there.\\n\\n\\n== Taxonomy ==\\nThe Southern African cheetah was first described by German naturalist Johann Christian Daniel von Schreber in his book Die Säugethiere in Abbildungen nach der Natur mit Beschreibungen (The Mammals illustrated as in Nature with Descriptions), published in 1775. Schreber described the species on basis of a specimen from the Cape of Good Hope. It is therefore the nominate subspecies. Subpopulations have been called \"South African cheetah\" and \"Namibian cheetah.\"Following Schreber\\'s description, other naturalists and zoologists also described cheetah specimens from many parts of Southern and East Africa that today are all considered synonyms of A. j. jubatus:\\nFelis guttata proposed in 1804 by Johann Hermann;\\nFelis fearonii proposed in 1834 by Andrew Smith;\\nFelis lanea proposed in 1877 by Philip Sclater;\\nAcinonyx jubatus obergi proposed in 1913 by Max Hilzheimer;\\nAcinonyx jubatus ngorongorensis proposed in 1913 by Hilzheimer on basis of a specimen from Ngorongoro, German East Africa;\\nAcinonyx jubatus velox proposed in 1913 by Edmund Heller on basis of a cheetah that was shot by Kermit Roosevelt in June 1909 in the Kenyan highlands.\\nAcinonyx rex proposed in 1927 by Reginald Innes Pocock on basis of a specimen from the Umvukwe Range in Rhodesia.In 2005, the authors of Mammal Species of the World grouped A. j. guttata, A. j. lanea, A. j. obergi, and A. j. rex under A j. jubatus, whilst recognizing A. j. raineyi and A. j. velox as valid taxa and considering P. l. ngorongore', metadata={'title': 'Southeast African cheetah', 'summary': 'The Southeast African cheetah (Acinonyx jubatus jubatus) is the nominate cheetah subspecies native to East and Southern Africa. The Southern African cheetah lives mainly in the lowland areas and deserts of the Kalahari, the savannahs of Okavango Delta, and the grasslands of the Transvaal region in South Africa. In Namibia, cheetahs are mostly found in farmlands. In India, four cheetahs of the subspecies are living in Kuno National Park in Madhya Pradesh after having been introduced there.', 'source': 'https://en.wikipedia.org/wiki/Southeast_African_cheetah'}),\n",
       "  Document(page_content='Footspeed, or sprint speed, is the maximum speed at which a human can run. It is affected by many factors, varies greatly throughout the population, and is important in athletics and many sports, such as association football, rugby football, American football,  track and field, field hockey, tennis, baseball, and basketball.\\n\\n\\n== Factors in speed ==\\nThe key determinant of footspeed in sprinting is the predominance of one distinct type of muscle fibre over another, specifically the ratio of fast-twitch muscles to slow-twitch muscles in a sprinter\\'s physical makeup.  Though fast-twitch muscles produce no more energy than slow-twitch muscles when they contract, they do so more rapidly through a process of anaerobic metabolism, though at the cost of inferior efficiency over longer periods of firing.  The average human has an almost-equal ratio of fast-twitch to slow-twitch fibers, but top sprinters may have as much as 80% fast-twitch fibers, while top long-distance runners may have only 20%. This ratio is believed to have genetic origins, though some assert that it can be adjusted by muscle training.  \"Speed camps\" and \"Speed Training Manuals\", which purport to provide fractional increases in maximum footspeed, are popular among budding professional athletes, and some sources estimate that 17–19% of speed can be trained.Though good running form is useful in increasing speed, fast and slow runners have been shown to move their legs at nearly the same rate – it is the force exerted by the leg on the ground that separates fast sprinters from slow.  Top short-distance runners exert as much as four times their body weight in pressure on the running surface. For this reason, muscle mass in the legs, relative to total body weight, is a key factor in maximizing footspeed.\\n\\n\\n== Limits of speed ==\\nThe record is 44.72 km/h (27.78 mph), measured between meter 60 and meter 80 of the 100 meters sprint at the 2009 World Championships in Athletics by Usain Bolt. (Bolt\\'s average speed o', metadata={'title': 'Footspeed', 'summary': 'Footspeed, or sprint speed, is the maximum speed at which a human can run. It is affected by many factors, varies greatly throughout the population, and is important in athletics and many sports, such as association football, rugby football, American football,  track and field, field hockey, tennis, baseball, and basketball.\\n\\n', 'source': 'https://en.wikipedia.org/wiki/Footspeed'}),\n",
       "  Document(page_content=\"This is a list of the fastest animals in the world, by types of animal.\\n\\n\\n== Fastest organism ==\\nThe peregrine falcon is the fastest bird, and the fastest member of the animal kingdom, with a diving speed of over 300 km/h (190 mph). The fastest land animal is the cheetah. Among the fastest animals in the sea is the black marlin, with uncertain and conflicting reports of recorded speeds.When drawing comparisons between different classes of animals, an alternative unit is sometimes used for organisms: body length per second. On this basis the 'fastest' organism on earth, relative to its body length, is the Southern Californian mite, Paratarsotomus macropalpis, which has a speed of 322 body lengths per second. The equivalent speed for a human, running as fast as this mite, would be 1,300 mph (2,092 km/h), or approximately Mach 1.7. The speed of the P. macropalpis is far in excess of the previous record holder, the Australian tiger beetle Cicindela eburneola, which is the fastest insect in the world relative to body size, with a recorded speed of 1.86 metres per second (6.7 km/h; 4.2 mph), or 171 body lengths per second. The cheetah, the fastest land mammal, scores at only 16 body lengths per second, while Anna's hummingbird has the highest known length-specific velocity attained by any vertebrate.\\n\\n\\n== Invertebrates ==\\n\\n\\n== Fish ==\\nDue to physical constraints, fish may be incapable of exceeding swim speeds of 36 km/h (22 mph). The larger reported figures below are therefore highly questionable:\\n\\n\\n== Amphibians ==\\n\\n\\n== Reptiles ==\\n\\n\\n== Birds ==\\n\\n\\n== Mammals ==\\n\\n\\n== See also ==\\nSpeed records\\n\\n\\n== Notes ==\\n\\n\\n== References ==\", metadata={'title': 'Fastest animals', 'summary': 'This is a list of the fastest animals in the world, by types of animal.', 'source': 'https://en.wikipedia.org/wiki/Fastest_animals'}),\n",
       "  Document(page_content=\"Pursuit predation is a form of predation in which predators actively give chase to their prey, either solitarily or as a group.  It is an alternate predation strategy to ambush predation — pursuit predators rely on superior speed, endurance and/or teamwork to seize the prey, while ambush predators use concealment, luring, exploiting of surroundings and the element of surprise to capture the prey.  While the two patterns of predation are not mutually exclusive, morphological differences in an organism's body plan can create an evolutionary bias favoring either type of predation.\\nPursuit predation is typically observed in carnivorous species within the kingdom Animalia, such as cheetahs, lions, wolves and early Homo species.  The chase can be initiated either by the predator, or by the prey if it is alerted to a predator's presence and attempt to flee before the predator gets close.  The chase ends either when the predator successfully catches up and tackles the prey, or when the predator abandons the attempt after the prey outruns it and escapes.\\nOne particular form of pursuit predation is persistence hunting, where the predator stalks the prey slowly but persistently to wear it down physically with fatigue or overheating; some animals are examples of both types of pursuit.\\n\\n\\n== Strategy ==\\nThere is still uncertainty as to whether predators behave with a general tactic or strategy while preying. However, among pursuit predators there are several common behaviors. Often, predators will scout potential prey, assessing prey quantity and density prior to engaging in a pursuit. Certain predators choose to pursue prey primarily in a group of conspecifics; these animals are known as pack hunters or group pursuers. Other species choose to hunt alone. These two behaviors are typically due to differences in hunting success, where some groups are very successful in groups and others are more successful alone. Pursuit predators may also choose to either exhaust their metabolic r\", metadata={'title': 'Pursuit predation', 'summary': \"Pursuit predation is a form of predation in which predators actively give chase to their prey, either solitarily or as a group.  It is an alternate predation strategy to ambush predation — pursuit predators rely on superior speed, endurance and/or teamwork to seize the prey, while ambush predators use concealment, luring, exploiting of surroundings and the element of surprise to capture the prey.  While the two patterns of predation are not mutually exclusive, morphological differences in an organism's body plan can create an evolutionary bias favoring either type of predation.\\nPursuit predation is typically observed in carnivorous species within the kingdom Animalia, such as cheetahs, lions, wolves and early Homo species.  The chase can be initiated either by the predator, or by the prey if it is alerted to a predator's presence and attempt to flee before the predator gets close.  The chase ends either when the predator successfully catches up and tackles the prey, or when the predator abandons the attempt after the prey outruns it and escapes.\\nOne particular form of pursuit predation is persistence hunting, where the predator stalks the prey slowly but persistently to wear it down physically with fatigue or overheating; some animals are examples of both types of pursuit.\", 'source': 'https://en.wikipedia.org/wiki/Pursuit_predation'})],\n",
       " 'annotations': [{'source_id': 0,\n",
       "   'quote': 'The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail.'}]}"
      ]
     },
     "execution_count": 92,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain_5.invoke(\"How fast are cheetahs?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "803c6155-48af-40db-b4b0-1ecc5328e99b",
   "metadata": {},
   "source": [
    "LangSmith trace: https://smith.langchain.com/public/8f30dbe5-9364-420c-9d90-63859ad06dcb/r\n",
    "\n",
    "If the answer was long we could first split it up and then apply the citation chain to every few sentences of the answer."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "poetry-venv",
   "language": "python",
   "name": "poetry-venv"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
